Merge pull request #4 from erikwilson/serve-context

Add context to serve
rancher · Aug 22, 2019 · 3d086ba · 3d086ba
2 parents 7de8789 + 7c71ffa
commit 3d086ba
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 12 deletions.
diff --git a/client.go b/client.go
@@ -18,7 +18,7 @@ func ClientConnect(ctx context.Context, wsURL string, headers http.Header, diale
 	}
 }
 
-func connectToProxy(ctx context.Context, proxyURL string, headers http.Header, auth ConnectAuthorizer, dialer *websocket.Dialer, onConnect func(context.Context) error) error {
+func connectToProxy(rootCtx context.Context, proxyURL string, headers http.Header, auth ConnectAuthorizer, dialer *websocket.Dialer, onConnect func(context.Context) error) error {
 	logrus.WithField("url", proxyURL).Info("Connecting to proxy")
 
 	if dialer == nil {
@@ -31,11 +31,11 @@ func connectToProxy(ctx context.Context, proxyURL string, headers http.Header, a
 	}
 	defer ws.Close()
 
-	if onConnect != nil {
-		ctxOnConnect, cancel := context.WithCancel(context.Background())
-		defer cancel()
+	ctx, cancel := context.WithCancel(rootCtx)
+	defer cancel()
 
-		if err := onConnect(ctxOnConnect); err != nil {
+	if onConnect != nil {
+		if err := onConnect(ctx); err != nil {
 			return err
 		}
 	}
@@ -45,7 +45,7 @@ func connectToProxy(ctx context.Context, proxyURL string, headers http.Header, a
 
 	result := make(chan error, 1)
 	go func() {
-		_, err = session.Serve()
+		_, err = session.Serve(ctx)
 		result <- err
 	}()
 

diff --git a/peer.go b/peer.go
@@ -106,7 +106,7 @@ outer:
 		}
 
 		s.sessions.addListener(session)
-		_, err = session.Serve()
+		_, err = session.Serve(context.Background())
 		s.sessions.removeListener(session)
 		session.Close()
 

diff --git a/server.go b/server.go
@@ -1,6 +1,7 @@
 package remotedialer
 
 import (
+	"context"
 	"net/http"
 	"sync"
 	"time"
@@ -71,7 +72,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	defer s.sessions.remove(session)
 
 	// Don't need to associate req.Context() to the Session, it will cancel otherwise
-	code, err := session.Serve()
+	code, err := session.Serve(context.Background())
 	if err != nil {
 		// Hijacked so we can't write to the client
 		logrus.Infof("error in remotedialer server [%d]: %v", code, err)

diff --git a/session.go b/session.go
@@ -63,8 +63,8 @@ func newSession(sessionKey int64, clientKey string, conn *websocket.Conn) *Sessi
 	}
 }
 
-func (s *Session) startPings() {
-	ctx, cancel := context.WithCancel(context.Background())
+func (s *Session) startPings(rootCtx context.Context) {
+	ctx, cancel := context.WithCancel(rootCtx)
 	s.pingCancel = cancel
 	s.pingWait.Add(1)
 
@@ -99,9 +99,9 @@ func (s *Session) stopPings() {
 	s.pingWait.Wait()
 }
 
-func (s *Session) Serve() (int, error) {
+func (s *Session) Serve(ctx context.Context) (int, error) {
 	if s.client {
-		s.startPings()
+		s.startPings(ctx)
 	}
 
 	for {