Respect context when writing to LocalServerReadyChan (#50)

Author: int128

e2e_test/context_test.go

@@ -13,7 +13,7 @@ import (
 	"golang.org/x/oauth2"
 )
 
-func TestContextCancel(t *testing.T) {
+func TestContextCancelOnWaitingForBrowser(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.TODO(), 100*time.Millisecond)
 	defer cancel()
 	s := httptest.NewServer(&authserver.Handler{
@@ -47,3 +47,41 @@ func TestContextCancel(t *testing.T) {
 		t.Errorf("err wants DeadlineExceeded but %+v", err)
 	}
 }
+
+func TestContextCancelOnLocalServerReadyChan(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.TODO(), 100*time.Millisecond)
+	defer cancel()
+	openBrowserCh := make(chan string)
+	defer close(openBrowserCh)
+	s := httptest.NewServer(&authserver.Handler{
+		T: t,
+		NewAuthorizationResponse: func(r authserver.AuthorizationRequest) string {
+			return fmt.Sprintf("%s?error=server_error", r.RedirectURI)
+		},
+		NewTokenResponse: func(r authserver.TokenRequest) (int, string) {
+			return 500, "should not reach here"
+		},
+	})
+	defer s.Close()
+	cfg := oauth2cli.Config{
+		OAuth2Config: oauth2.Config{
+			ClientID:     "YOUR_CLIENT_ID",
+			ClientSecret: "YOUR_CLIENT_SECRET",
+			Scopes:       []string{"email", "profile"},
+			Endpoint: oauth2.Endpoint{
+				AuthURL:  s.URL + "/auth",
+				TokenURL: s.URL + "/token",
+			},
+		},
+		LocalServerReadyChan: openBrowserCh,
+		Logf:                 t.Logf,
+	}
+	_, err := oauth2cli.GetToken(ctx, cfg)
+	if err == nil {
+		t.Errorf("GetToken wants error but was nil")
+		return
+	}
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Errorf("err wants DeadlineExceeded but %+v", err)
+	}
+}

server.go

@@ -66,10 +66,18 @@ func receiveCodeViaLocalServer(ctx context.Context, c *Config) (string, error) {
 		}
 		return nil
 	})
-	if c.LocalServerReadyChan != nil {
-		c.LocalServerReadyChan <- c.OAuth2Config.RedirectURL
-	}
-
+	eg.Go(func() error {
+		if c.LocalServerReadyChan == nil {
+			return nil
+		}
+		select {
+		case c.LocalServerReadyChan <- c.OAuth2Config.RedirectURL:
+			c.Logf("oauth2cli: wrote a URL to LocalServerReadyChan: %s", c.OAuth2Config.RedirectURL)
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	})
 	if err := eg.Wait(); err != nil {
 		return "", fmt.Errorf("authorization error: %w", err)
 	}