Move ResponseWriter Conn and TLS methods to Request

3 files changed, 19 insertions, 21 deletions

diff --git a/request.go b/request.go
index 3bfa3db..f15d891 100644
--- a/request.go
+++ b/request.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"crypto/tls"
 	"io"
+	"net"
 	"net/url"
 )
 
@@ -26,6 +27,8 @@ type Request struct {
 	// TLS certificate to present to the other side of the connection.
 	// This field is ignored by the Gemini server.
 	Certificate *tls.Certificate
+
+	conn net.Conn
 }
 
 // NewRequest returns a new request.
@@ -89,3 +92,18 @@ func (r *Request) Write(w io.Writer) error {
 	}
 	return bw.Flush()
 }
+
+// Conn returns the network connection on which the request was received.
+func (r *Request) Conn() net.Conn {
+	return r.conn
+}
+
+// TLS returns information about the TLS connection on which the
+// request was received.
+func (r *Request) TLS() *tls.ConnectionState {
+	if tlsConn, ok := r.conn.(*tls.Conn); ok {
+		state := tlsConn.ConnectionState()
+		return &state
+	}
+	return nil
+}
diff --git a/response.go b/response.go
index 487e58b..d9d7503 100644
--- a/response.go
+++ b/response.go
@@ -156,13 +156,6 @@ type ResponseWriter interface {
 	// Any blocked Write operations will be unblocked and return errors.
 	Close() error
 
-	// Conn returns the underlying network connection.
-	// To take over the connection, use Hijack.
-	Conn() net.Conn
-
-	// TLS returns information about the underlying TLS connection.
-	TLS() *tls.ConnectionState
-
 	// unexported method so we can extend this interface over time
 	// without breaking existing code. Implementers must embed a concrete
 	// type from elsewhere.
@@ -175,7 +168,6 @@ type responseWriter struct {
 	mediatype   string
 	wroteHeader bool
 	bodyAllowed bool
-	conn        net.Conn
 }
 
 func newResponseWriter(w io.WriteCloser) *responseWriter {
@@ -232,16 +224,4 @@ func (w *responseWriter) Close() error {
 	return w.cl.Close()
 }
 
-func (w *responseWriter) Conn() net.Conn {
-	return w.conn
-}
-
-func (w *responseWriter) TLS() *tls.ConnectionState {
-	if tlsConn, ok := w.conn.(*tls.Conn); ok {
-		state := tlsConn.ConnectionState()
-		return &state
-	}
-	return nil
-}
-
 func (w *responseWriter) unexported() {}
diff --git a/server.go b/server.go
index 9495caa..fb8f287 100644
--- a/server.go
+++ b/server.go
@@ -360,13 +360,13 @@ func (srv *Server) serveConn(ctx context.Context, conn net.Conn) error {
 	}
 
 	w := newResponseWriter(cw)
-	w.conn = conn
 
 	req, err := ReadRequest(r)
 	if err != nil {
 		w.WriteHeader(StatusBadRequest, "Bad request")
 		return w.Flush()
 	}
+	req.conn = conn
 
 	h := srv.Handler
 	if h == nil {