Merge context_conn into request_context

Summary: Put them in one file.

Reviewed By: podtserkovskiy

Differential Revision: D95894927

fbshipit-source-id: 43cb21661038d070fa1cebacc43158279ce64c2c

Author: echistyakov

thrift/lib/go/thrift/context_conn.go

@@ -18,17 +18,66 @@ package thrift
 
 import (
 	"context"
+	"crypto/tls"
 	"maps"
 	"net"
 	"time"
 )
 
+type contextKey int
+
+const (
+	connInfoKey   contextKey = 0
+	reqContextKey contextKey = 2
+)
+
 // Identity represents a secure peer identity
 type Identity struct {
 	Type string
 	Data string
 }
 
+// ConnInfo contains connection information from clients of the Server.
+type ConnInfo struct {
+	RemoteAddr net.Addr
+	tlsState   tlsConnectionStater // set by thrift tcp servers
+}
+
+// tlsConnectionStater is an abstract interface for types that can return
+// the state of TLS connections. This is used to support not only tls.Conn
+// but also custom wrappers such as permissive TLS/non-TLS sockets.
+//
+// Caveat: this interface has to support at least tls.Conn, which has
+// the current signature for ConnectionState. Because of that, wrappers
+// for permissive TLS/non-TLS may return an empty tls.ConnectionState.
+type tlsConnectionStater interface {
+	ConnectionState() tls.ConnectionState
+}
+
+// tlsConnectionStaterHandshaker is an abstract interface that allows
+// custom "TLS-like" connections to be used with Thrift ALPN logic.
+type tlsConnectionStaterHandshaker interface {
+	tlsConnectionStater
+	HandshakeContext(context.Context) error
+}
+
+// Compile time interface enforcer
+var _ tlsConnectionStater = (*tls.Conn)(nil)
+var _ tlsConnectionStaterHandshaker = (*tls.Conn)(nil)
+
+// TLS returns the TLS connection state.
+func (c ConnInfo) TLS() *tls.ConnectionState {
+	if c.tlsState == nil {
+		return nil
+	}
+	cs := c.tlsState.ConnectionState()
+	// See the caveat in tlsConnectionStater.
+	if cs.Version == 0 {
+		return nil
+	}
+	return &cs
+}
+
 // RequestContext is a mirror of C++ apache::thrift::RequestContext
 // Not all options are guaranteed to be implemented by a client
 type RequestContext struct {
@@ -50,11 +99,24 @@ type RequestContext struct {
 	contextHeaders
 }
 
-type requestContextKey int
+// withConnInfo adds connection info (from a thrift.Transport) to context, if applicable
+func withConnInfo(ctx context.Context, conn net.Conn) context.Context {
+	var tlsState tlsConnectionStater
+	if t, ok := conn.(tlsConnectionStater); ok {
+		tlsState = t
+	}
+	ctx = context.WithValue(ctx, connInfoKey, ConnInfo{
+		RemoteAddr: conn.RemoteAddr(),
+		tlsState:   tlsState,
+	})
+	return ctx
+}
 
-const (
-	reqContextKey requestContextKey = 2
-)
+// connInfoFromContext extracts and returns ConnInfo from context.
+func connInfoFromContext(ctx context.Context) (ConnInfo, bool) {
+	v, ok := ctx.Value(connInfoKey).(ConnInfo)
+	return v, ok
+}
 
 // GetRequestContext returns the RequestContext in a go context, or nil if there is nothing
 func GetRequestContext(ctx context.Context) *RequestContext {