Initial import

proxy/session.go

@@ -0,0 +1,151 @@
+package proxy
+
+import (
+	"bufio"
+	"crypto/tls"
+	"encoding/binary"
+	"encoding/hex"
+	"math/rand"
+	"net"
+	"net/http"
+	"sync/atomic"
+	"time"
+
+	"git.maze.io/maze/styx/internal/log"
+)
+
+var seed = rand.NewSource(time.Now().UnixNano())
+
+type Context struct {
+	id     int64
+	conn   *wrappedConn
+	rw     *bufio.ReadWriter
+	parent *Session
+	data   map[string]any
+}
+
+func newContext(conn net.Conn, rw *bufio.ReadWriter, parent *Session) *Context {
+	if wrapped, ok := conn.(*wrappedConn); ok {
+		conn = wrapped.Conn
+	}
+
+	ctx := &Context{
+		id:     seed.Int63(),
+		conn:   &wrappedConn{Conn: conn},
+		rw:     rw,
+		parent: parent,
+		data:   make(map[string]any),
+	}
+
+	return ctx
+}
+
+func (ctx *Context) log() log.Logger {
+	return log.Console.With().
+		Str("context", ctx.ID()).
+		Str("addr", ctx.RemoteAddr().String()).
+		Logger()
+}
+
+func (ctx *Context) ID() string {
+	var b [8]byte
+	binary.BigEndian.PutUint64(b[:], uint64(ctx.id))
+	if ctx.parent != nil {
+		return ctx.parent.ID() + "-" + hex.EncodeToString(b[:])
+	}
+	return hex.EncodeToString(b[:])
+}
+
+func (ctx *Context) IsTLS() bool {
+	_, ok := ctx.conn.Conn.(*tls.Conn)
+	return ok && ctx.parent != nil
+}
+
+func (ctx *Context) RemoteAddr() net.Addr {
+	if ctx.parent != nil {
+		return ctx.parent.ctx.RemoteAddr()
+	}
+	return ctx.conn.RemoteAddr()
+}
+
+func (ctx *Context) SetDeadline(t time.Time) error {
+	if ctx.parent != nil {
+		return ctx.parent.ctx.SetDeadline(t)
+	}
+	return ctx.conn.SetDeadline(t)
+}
+
+func (ctx *Context) Set(key string, value any) {
+	ctx.data[key] = value
+}
+
+func (ctx *Context) Get(key string) (value any, ok bool) {
+	value, ok = ctx.data[key]
+	return
+}
+
+func (ctx *Context) Flush() error {
+	return ctx.rw.Flush()
+}
+
+func (ctx *Context) Write(p []byte) (n int, err error) {
+	if n, err = ctx.rw.Write(p); n > 0 {
+		atomic.AddInt64(&ctx.conn.bytes, int64(n))
+	}
+	return
+}
+
+type Session struct {
+	id       int64
+	ctx      *Context
+	request  *http.Request
+	response *http.Response
+	data     map[string]any
+}
+
+func newSession(ctx *Context, request *http.Request) *Session {
+	return &Session{
+		id:      seed.Int63(),
+		ctx:     ctx,
+		request: request,
+		data:    make(map[string]any),
+	}
+}
+
+func (ses *Session) log() log.Logger {
+	return log.Console.With().
+		Str("context", ses.ctx.ID()).
+		Str("session", ses.ID()).
+		Str("addr", ses.ctx.RemoteAddr().String()).
+		Logger()
+}
+
+func (ses *Session) ID() string {
+	var b [8]byte
+	binary.BigEndian.PutUint64(b[:], uint64(ses.id))
+	return hex.EncodeToString(b[:])
+}
+
+func (ses *Session) Context() *Context {
+	return ses.ctx
+}
+
+func (ses *Session) Request() *http.Request {
+	return ses.request
+}
+
+func (ses *Session) Response() *http.Response {
+	return ses.response
+}
+
+type wrappedConn struct {
+	net.Conn
+	bytes int64
+}
+
+func (c *wrappedConn) Write(p []byte) (n int, err error) {
+	if n, err = c.Conn.Write(p); n > 0 {
+		atomic.AddInt64(&c.bytes, int64(n))
+	}
+	return
+}