Better trie implementations

proxy/context.go

@@ -45,8 +45,11 @@ type Context interface {
 	// Response is the response that will be sent back to the client.
 	Response() *http.Response
 
+	// Logger for this context.
+	Logger() logger.Structured
+
 	// Client group.
-	Client() (dataset.Client, error)
+	Storage() dataset.Storage
 }
 
 type WithCertificateAuthority interface {
@@ -91,11 +94,10 @@ type proxyContext struct {
 	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 {
+func NewContext(c net.Conn, storage dataset.Storage) Context {
 	if c, ok := c.(*proxyContext); ok {
 		return c
 	}
@@ -106,12 +108,13 @@ func NewContext(c net.Conn) Context {
 	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},
+		Conn:    c,
+		id:      binary.BigEndian.Uint64(b),
+		cr:      cr,
+		br:      bufio.NewReader(cr),
+		cw:      cw,
+		res:     &http.Response{StatusCode: 200},
+		storage: storage,
 	}
 }
 
@@ -128,7 +131,7 @@ func (c *proxyContext) AccessLogEntry() logger.Structured {
 	return entry
 }
 
-func (c *proxyContext) LogEntry() logger.Structured {
+func (c *proxyContext) Logger() logger.Structured {
 	var id [8]byte
 	binary.BigEndian.PutUint64(id[:], c.id)
 	return ServerLog.Values(logger.Values{
@@ -234,24 +237,8 @@ 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
+func (c *proxyContext) Storage() dataset.Storage {
+	return c.storage
 }
 
 var _ Context = (*proxyContext)(nil)