whoami command just about working

authorize/ssh_grpc.go

@@ -1,6 +1,7 @@
 package authorize
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"errors"
@@ -10,11 +11,13 @@ import (
 	"slices"
 	"strings"
 	"sync/atomic"
+	"text/template"
 	"time"
 
 	"github.com/charmbracelet/bubbles/list"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
+	envoy_config_core_v3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extensions_ssh "github.com/pomerium/envoy-custom/api/extensions/filters/network/ssh"
 	"github.com/pomerium/pomerium/authorize/evaluator"
 	"github.com/pomerium/pomerium/config"
@@ -28,8 +31,10 @@ import (
 	"golang.org/x/oauth2"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
@@ -411,6 +416,9 @@ func handleEvaluatorResponseForSSH(
 	// TODO: ideally there would be a way to keep this in sync with the logic in check_response.go
 	allow := result.Allow.Value && !result.Deny.Value
 
+	// XXX
+	sessionID, _ := getSessionIDForSSH(state.PublicKey)
+
 	if allow {
 		pkData, _ := anypb.New(publicKeyAllowResponse(state.PublicKey))
 		return &extensions_ssh.ServerMessage{
@@ -431,6 +439,15 @@ func handleEvaluatorResponseForSSH(
 							},
 							//Target: extensions_ssh.Target_Upstream,
 							Target: extensions_ssh.Target_Internal,
+							SetMetadata: &envoy_config_core_v3.Metadata{
+								FilterMetadata: map[string]*structpb.Struct{
+									"pomerium.ssh": {
+										Fields: map[string]*structpb.Value{
+											"pomerium-session-id": structpb.NewStringValue(sessionID),
+										},
+									},
+								},
+							},
 						},
 					},
 				},
@@ -607,6 +624,7 @@ func (a *Authorize) ServeChannel(
 				},
 			},
 		})
+		fmt.Println(" *** sending handoff request *** ")
 		return server.Send(&extensions_ssh.ChannelMessage{
 			Message: &extensions_ssh.ChannelMessage_ChannelControl{
 				ChannelControl: &extensions_ssh.ChannelControl{
@@ -662,33 +680,12 @@ func (a *Authorize) ServeChannel(
 			fmt.Println(" *** SSH_MSG_CHANNEL_REQUEST: ", msg.Request)
 
 			switch msg.Request {
-			case "pty-req":
+			case "env":
-				req := parsePtyReq(msg.RequestSpecificData)
+				// ignore for now
-				downstreamPtyInfo = &extensions_ssh.SSHDownstreamPTYInfo{
-					TermEnv:      req.TermEnv,
-					WidthColumns: req.Width,
-					HeightRows:   req.Height,
-					WidthPx:      req.WidthPx,
-					HeightPx:     req.HeightPx,
-					Modes:        req.Modes,
-				}
-				if err := server.Send(&extensions_ssh.ChannelMessage{
-					Message: &extensions_ssh.ChannelMessage_RawBytes{
-						RawBytes: &wrapperspb.BytesValue{
-							Value: gossh.Marshal(channelRequestSuccessMsg{
-								PeersID: peerId,
-							}),
-						},
-					},
-				}); err != nil {
-					return err
-				}
 			case "subsystem":
 				subsystem := parseString(msg.RequestSpecificData)
-				fmt.Println("     -> subsystem: ", subsystem)
+				command, isInternal := strings.CutPrefix(subsystem, "pomerium ")
-				switch subsystem {
+				if isInternal {
-				case "pomerium-whoami":
-					fmt.Println(" *** who am I? ***")
 					if err := server.Send(&extensions_ssh.ChannelMessage{
 						Message: &extensions_ssh.ChannelMessage_RawBytes{
 							RawBytes: &wrapperspb.BytesValue{
@@ -700,24 +697,9 @@ func (a *Authorize) ServeChannel(
 					}); err != nil {
 						return err
 					}
-					if err := server.Send(&extensions_ssh.ChannelMessage{
+					return a.serveInternalCommand(server, peerId, command)
-						Message: &extensions_ssh.ChannelMessage_RawBytes{
+				} else if err := handoff(); err != nil {
-							RawBytes: &wrapperspb.BytesValue{
+					return err
-								Value: gossh.Marshal(channelDataMsg{
-									PeersID: peerId,
-									Length:  uint32(12),
-									Rest:    []byte("hello world!"),
-								}),
-							},
-						},
-					}); err != nil {
-						return err
-					}
-					return nil // close the stream
-				default:
-					if err := handoff(); err != nil {
-						return err
-					}
 				}
 			default:
 				// We're not interested in hijacking any other kinds of session.
@@ -910,6 +892,62 @@ func (a *Authorize) ServeChannel(
 	}
 }*/
 
+var whoamiTmpl = template.Must(template.New("whoami").Parse(`
+User ID:    {{.UserId}}
+Session ID: {{.Id}}
+Expires at: {{.ExpiresAt.AsTime}}
+Claims:
+{{- range $k, $v := .Claims }}
+  {{ $k }}: {{ $v.AsSlice }}
+{{- end }}
+`))
+
+func (a *Authorize) serveInternalCommand(
+	server extensions_ssh.StreamManagement_ServeChannelServer,
+	peerID uint32,
+	command string,
+) error {
+	md, ok := metadata.FromIncomingContext(server.Context())
+	fmt.Println("metadata.FromIncomingContext: ", md, ok)
+
+	var sessionID string
+	if h := md.Get("pomerium-session-id"); len(h) == 1 {
+		sessionID = h[0]
+	}
+
+	switch command {
+	case "whoami":
+		fmt.Println(" *** who am I ? ***")
+		client := a.state.Load().dataBrokerClient
+		var output string
+		s, err := session.Get(server.Context(), client, sessionID)
+		if err != nil {
+			output = fmt.Sprint("couldn't fetch session: ", err.Error())
+		} else {
+			var b bytes.Buffer
+			whoamiTmpl.Execute(&b, s)
+			output = b.String()
+		}
+
+		// TODO
+		if err := server.Send(&extensions_ssh.ChannelMessage{
+			Message: &extensions_ssh.ChannelMessage_RawBytes{
+				RawBytes: &wrapperspb.BytesValue{
+					Value: gossh.Marshal(channelDataMsg{
+						PeersID: peerID,
+						Length:  uint32(len(output)),
+						Rest:    []byte(output),
+					}),
+				},
+			},
+		}); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 type ptyReq struct {
 	TermEnv           string
 	Width, Height     uint32

go.mod

@@ -16,8 +16,6 @@ require (
 	github.com/charmbracelet/bubbles v0.20.0
 	github.com/charmbracelet/bubbletea v1.3.3
 	github.com/charmbracelet/lipgloss v1.0.0
-	github.com/charmbracelet/x/ansi v0.8.0
-	github.com/charmbracelet/x/term v0.2.1
 	github.com/cloudflare/circl v1.5.0
 	github.com/cncf/xds/go v0.0.0-20241223141626-cff3c89139a3
 	github.com/coreos/go-oidc/v3 v3.11.0
@@ -144,6 +142,8 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/caddyserver/zerossl v0.1.3 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/charmbracelet/x/ansi v0.8.0 // indirect
+	github.com/charmbracelet/x/term v0.2.1 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/containerd/platforms v0.2.1 // indirect
 	github.com/cpuguy83/dockercfg v0.3.2 // indirect