fix: context data race (#4)

Author: aymanbagabas

context.go

@@ -93,12 +93,12 @@ type Context interface {
 
 type sshContext struct {
 	context.Context
-	*sync.Mutex
+	*sync.RWMutex
 }
 
 func newContext(srv *Server) (*sshContext, context.CancelFunc) {
 	innerCtx, cancel := context.WithCancel(context.Background())
-	ctx := &sshContext{innerCtx, &sync.Mutex{}}
+	ctx := &sshContext{innerCtx, &sync.RWMutex{}}
 	ctx.SetValue(ContextKeyServer, srv)
 	perms := &Permissions{&gossh.Permissions{}}
 	ctx.SetValue(ContextKeyPermissions, perms)
@@ -120,9 +120,23 @@ func applyConnMetadata(ctx Context, conn gossh.ConnMetadata) {
 }
 
 func (ctx *sshContext) SetValue(key, value interface{}) {
+	ctx.RWMutex.Lock()
+	defer ctx.RWMutex.Unlock()
 	ctx.Context = context.WithValue(ctx.Context, key, value)
 }
 
+func (ctx *sshContext) Value(key interface{}) interface{} {
+	ctx.RWMutex.RLock()
+	defer ctx.RWMutex.RUnlock()
+	return ctx.Context.Value(key)
+}
+
+func (ctx *sshContext) Done() <-chan struct{} {
+	ctx.RWMutex.RLock()
+	defer ctx.RWMutex.RUnlock()
+	return ctx.Context.Done()
+}
+
 func (ctx *sshContext) User() string {
 	return ctx.Value(ContextKeyUser).(string)
 }

context_test.go

@@ -45,3 +45,27 @@ func TestSetValue(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestRaceRWIssue160(t *testing.T) {
+	value := "foo"
+	key := "bar"
+	session, _, cleanup := newTestSessionWithOptions(t, &Server{
+		Handler: func(s Session) {
+			t.Run("test done", func(t *testing.T) {
+				t.Parallel()
+				go func() {
+					s.Context().SetValue(key, value)
+				}()
+				go func() {
+					select {
+					case <-s.Context().Done():
+					}
+				}()
+			})
+		},
+	}, nil)
+	defer cleanup()
+	if err := session.Run(""); err != nil {
+		t.Fatal(err)
+	}
+}