repos / pico

commit

0a1449a

parent

2f2b7cc

author

Antonio Mika

date

2025-03-12 23:31:32 -0400 EDT

Add session specific context based on conn context

1@@ -9,3 +9,4 @@ build/*
2 data/*
3 !data/.gitkeep
4 Dockerfile
5+ssh_data

  1@@ -35,6 +35,13 @@ type SSHServerConn struct {
  2 	mu sync.Mutex
  3 }
  4 
  5+func (s *SSHServerConn) Context() context.Context {
  6+	s.mu.Lock()
  7+	defer s.mu.Unlock()
  8+
  9+	return s.Ctx
 10+}
 11+
 12 func (sc *SSHServerConn) Close() error {
 13 	sc.CancelFunc()
 14 	return nil
 15@@ -44,6 +51,9 @@ type SSHServerConnSession struct {
 16 	ssh.Channel
 17 	*SSHServerConn
 18 
 19+	Ctx        context.Context
 20+	CancelFunc context.CancelFunc
 21+
 22 	pty   *Pty
 23 	winch chan Window
 24 
 25@@ -51,7 +61,7 @@ type SSHServerConnSession struct {
 26 }
 27 
 28 // Deadline implements context.Context.
 29-func (s *SSHServerConn) Deadline() (deadline time.Time, ok bool) {
 30+func (s *SSHServerConnSession) Deadline() (deadline time.Time, ok bool) {
 31 	s.mu.Lock()
 32 	defer s.mu.Unlock()
 33 
 34@@ -59,7 +69,7 @@ func (s *SSHServerConn) Deadline() (deadline time.Time, ok bool) {
 35 }
 36 
 37 // Done implements context.Context.
 38-func (s *SSHServerConn) Done() <-chan struct{} {
 39+func (s *SSHServerConnSession) Done() <-chan struct{} {
 40 	s.mu.Lock()
 41 	defer s.mu.Unlock()
 42 
 43@@ -67,7 +77,7 @@ func (s *SSHServerConn) Done() <-chan struct{} {
 44 }
 45 
 46 // Err implements context.Context.
 47-func (s *SSHServerConn) Err() error {
 48+func (s *SSHServerConnSession) Err() error {
 49 	s.mu.Lock()
 50 	defer s.mu.Unlock()
 51 
 52@@ -75,7 +85,7 @@ func (s *SSHServerConn) Err() error {
 53 }
 54 
 55 // Value implements context.Context.
 56-func (s *SSHServerConn) Value(key any) any {
 57+func (s *SSHServerConnSession) Value(key any) any {
 58 	s.mu.Lock()
 59 	defer s.mu.Unlock()
 60 
 61@@ -83,14 +93,14 @@ func (s *SSHServerConn) Value(key any) any {
 62 }
 63 
 64 // SetValue implements context.Context.
 65-func (s *SSHServerConn) SetValue(key any, data any) {
 66+func (s *SSHServerConnSession) SetValue(key any, data any) {
 67 	s.mu.Lock()
 68 	defer s.mu.Unlock()
 69 
 70 	s.Ctx = context.WithValue(s.Ctx, key, data)
 71 }
 72 
 73-func (s *SSHServerConn) Context() context.Context {
 74+func (s *SSHServerConnSession) Context() context.Context {
 75 	s.mu.Lock()
 76 	defer s.mu.Unlock()
 77 
 78@@ -128,6 +138,7 @@ func (s *SSHServerConnSession) Command() []string {
 79 }
 80 
 81 func (s *SSHServerConnSession) Close() error {
 82+	s.CancelFunc()
 83 	return s.Channel.Close()
 84 }
 85 
 86@@ -162,7 +173,7 @@ func (sc *SSHServerConn) Handle(chans <-chan ssh.NewChannel, reqs <-chan *ssh.Re
 87 
 88 	for {
 89 		select {
 90-		case <-sc.Done():
 91+		case <-sc.Context().Done():
 92 			return nil
 93 		case newChan, ok := <-chans:
 94 			if !ok {
 95@@ -399,14 +410,18 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon
 96 				return err
 97 			}
 98 
 99+			ctx, cancelFunc := context.WithCancel(sc.Ctx)
100+
101 			sesh := &SSHServerConnSession{
102 				Channel:       channel,
103 				SSHServerConn: sc,
104+				Ctx:           ctx,
105+				CancelFunc:    cancelFunc,
106 			}
107 
108 			for {
109 				select {
110-				case <-sc.Done():
111+				case <-sesh.Done():
112 					return nil
113 				case req, ok := <-requests:
114 					if !ok {