Making sure connection value is available in streams, as well as the …

…stream itself
loopholelabs · Sep 11, 2024 · a60dfc4 · a60dfc4
1 parent 5672a73
commit a60dfc4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/templates/constants.templ b/templates/constants.templ
@@ -1,5 +1,6 @@
 {{define "constants"}}
     const ConnectionContextKey int = 1000
+    const StreamContextKey int = 1001
 
     func SetErrorFlag(flags uint8, error bool) uint8 {
         return flags | 0x2

diff --git a/templates/server.templ b/templates/server.templ
@@ -54,6 +54,8 @@ func NewServer({{ GetServerFields .services }}, tlsConfig *tls.Config, logger ty
 
     {{ if .numStreamMethods -}}
     s.server.SetStreamHandler(func(ctx context.Context, stream *frisbee.Stream) {
+        streamCtx := context.WithValue(ctx, ConnectionContextKey, stream.Conn())
+        streamCtx = context.WithValue(streamCtx, StreamContextKey, stream)
         p, err := stream.ReadPacket()
         if err != nil {
             return
@@ -73,7 +75,7 @@ func NewServer({{ GetServerFields .services }}, tlsConfig *tls.Config, logger ty
                     {{ $opIndex := call $counter -}}
                     {{ if or $method.IsStreamingClient $method.IsStreamingServer -}}
                         case {{ $opIndex }}:
-                            s.create{{ CamelCaseName $method.Name }}Server(ctx, {{ FirstLowerCase (CamelCaseName $service.Name) }}, stream)
+                            s.create{{ CamelCaseName $method.Name }}Server(streamCtx, {{ FirstLowerCase (CamelCaseName $service.Name) }}, stream)
                     {{end -}}
                 {{end -}}
             {{end -}}