styx/proxy/context.go

package proxy

import (
	"bufio"
	"crypto/rand"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"sync/atomic"
	"time"

	"git.maze.io/maze/styx/ca"
	"git.maze.io/maze/styx/dataset"
	"git.maze.io/maze/styx/logger"
)

// Context provides convenience functions for the current ongoing HTTP proxy transaction (request).
type Context interface {
	// Conn is the backing connection for this context.
	net.Conn

	// ID is a unique connection identifier.
	ID() uint64

	// Reader returns a buffered reader on top of the [net.Conn].
	Reader() *bufio.Reader

	// BytesRead returns the number of bytes read.
	BytesRead() int64

	// BytesSent returns the number of bytes written.
	BytesSent() int64

	// TLSState returns the TLS connection state, it returns nil if the connection is not a TLS connection.
	TLSState() *tls.ConnectionState

	// Request is the request made to the proxy.
	Request() *http.Request

	// Response is the response that will be sent back to the client.
	Response() *http.Response

	// Client group.
	Client() (dataset.Client, error)
}

type WithCertificateAuthority interface {
	CertificateAuthority() ca.CertificateAuthority
}

type countingReader struct {
	reader io.Reader
	bytes  int64
}

func (r *countingReader) Read(p []byte) (n int, err error) {
	if n, err = r.reader.Read(p); n > 0 {
		atomic.AddInt64(&r.bytes, int64(n))
	}
	return
}

type countingWriter struct {
	writer io.Writer
	bytes  int64
}

func (w *countingWriter) Write(p []byte) (n int, err error) {
	if n, err = w.writer.Write(p); n > 0 {
		atomic.AddInt64(&w.bytes, int64(n))
	}
	return
}

type proxyContext struct {
	net.Conn
	id             uint64
	cr             *countingReader
	br             *bufio.Reader
	cw             *countingWriter
	transparent    int
	transparentTLS bool
	serverName     string
	req            *http.Request
	res            *http.Response
	idleTimeout    time.Duration
	ca             ca.CertificateAuthority
	storage        dataset.Storage
	client         dataset.Client
}

// NewContext returns an initialized context for the provided [net.Conn].
func NewContext(c net.Conn) Context {
	if c, ok := c.(*proxyContext); ok {
		return c
	}

	b := make([]byte, 8)
	io.ReadFull(rand.Reader, b)

	cr := &countingReader{reader: c}
	cw := &countingWriter{writer: c}
	return &proxyContext{
		Conn: c,
		id:   binary.BigEndian.Uint64(b),
		cr:   cr,
		br:   bufio.NewReader(cr),
		cw:   cw,
		res:  &http.Response{StatusCode: 200},
	}
}

func (c *proxyContext) AccessLogEntry() logger.Structured {
	var id [8]byte
	binary.BigEndian.PutUint64(id[:], c.id)
	entry := AccessLog.Values(logger.Values{
		"client":   c.RemoteAddr().String(),
		"server":   c.LocalAddr().String(),
		"id":       hex.EncodeToString(id[:]),
		"bytes_rx": c.BytesRead(),
		"bytes_tx": c.BytesSent(),
	})
	return entry
}

func (c *proxyContext) LogEntry() logger.Structured {
	var id [8]byte
	binary.BigEndian.PutUint64(id[:], c.id)
	return ServerLog.Values(logger.Values{
		"client": c.RemoteAddr().String(),
		"server": c.LocalAddr().String(),
		"id":     hex.EncodeToString(id[:]),
	})
}

func (c *proxyContext) String() string {
	return fmt.Sprintf("client=%s server=%s id=%#08x",
		c.RemoteAddr().String(),
		c.LocalAddr().String(),
		c.id)
}

func (c *proxyContext) ID() uint64 {
	return c.id
}

func (c *proxyContext) BytesRead() int64 {
	return atomic.LoadInt64(&c.cr.bytes)
}

func (c *proxyContext) BytesSent() int64 {
	return atomic.LoadInt64(&c.cw.bytes)
}

func (c *proxyContext) Read(p []byte) (n int, err error) {
	if c.idleTimeout > 0 {
		if err = c.SetReadDeadline(time.Now().Add(c.idleTimeout)); err != nil {
			return
		}
	}
	return c.br.Read(p)
}

func (c *proxyContext) Write(p []byte) (n int, err error) {
	if c.idleTimeout > 0 {
		if err = c.SetWriteDeadline(time.Now().Add(c.idleTimeout)); err != nil {
			return
		}
	}
	return c.cw.Write(p)
}

func (c *proxyContext) Reader() *bufio.Reader {
	return c.br
}

func (c *proxyContext) Request() *http.Request {
	return c.req
}

func (c *proxyContext) SetRequest(req *http.Request) {
	c.req = req
}

func (c *proxyContext) Response() *http.Response {
	return c.res
}

func (c *proxyContext) SetIdleTimeout(t time.Duration) {
	c.idleTimeout = t
}

type connectionStater interface {
	ConnectionState() tls.ConnectionState
}

func (c *proxyContext) TLSState() *tls.ConnectionState {
	if s, ok := c.Conn.(connectionStater); ok {
		state := s.ConnectionState()
		return &state
	}
	return nil
}

// http.ResponseWriter interface:

func (c *proxyContext) Header() http.Header {
	if c.res == nil {
		c.res = NewResponse(http.StatusOK, nil, c.req)
	}
	return c.res.Header
}

func (c *proxyContext) WriteHeader(code int) {
	if c.res == nil {
		c.res = NewResponse(code, nil, c.req)
	} else {
		if text := http.StatusText(code); text != "" {
			c.res.Status = strconv.Itoa(code) + " " + text
		} else {
			c.res.Status = strconv.Itoa(code)
		}
		c.res.StatusCode = code
	}
	//return c.res.Header.Write(c)
}

func (c *proxyContext) CertificateAuthority() ca.CertificateAuthority {
	return c.ca
}

func (c *proxyContext) Client() (dataset.Client, error) {
	if c.storage == nil {
		return dataset.Client{}, dataset.ErrNotExist{Object: "client"}
	}
	if !c.client.CreatedAt.Equal(time.Time{}) {
		return c.client, nil
	}

	var err error
	switch addr := c.Conn.RemoteAddr().(type) {
	case *net.TCPAddr:
		c.client, err = c.storage.ClientByIP(addr.IP)
	case *net.UDPAddr:
		c.client, err = c.storage.ClientByIP(addr.IP)
	default:
		err = dataset.ErrNotExist{Object: "client"}
	}
	return c.client, err
}

var _ Context = (*proxyContext)(nil)