152 lines
2.8 KiB
Go
152 lines
2.8 KiB
Go
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
|
|
}
|