plumbing: wire up contexts for Transport.AdvertisedReferences (#246)

* plumbing: wire up contexts for Transport.AdvertisedReferences

* add more tests for context wiring

7 files changed, 64 insertions, 7 deletions

diff --git a/plumbing/transport/common.go b/plumbing/transport/common.go
index b993c4e..a9ee2ca 100644
--- a/plumbing/transport/common.go
+++ b/plumbing/transport/common.go
@@ -58,6 +58,11 @@ type Session interface {
 	// If the repository does not exist, returns ErrRepositoryNotFound.
 	// If the repository exists, but is empty, returns ErrEmptyRemoteRepository.
 	AdvertisedReferences() (*packp.AdvRefs, error)
+	// AdvertisedReferencesContext retrieves the advertised references for a
+	// repository.
+	// If the repository does not exist, returns ErrRepositoryNotFound.
+	// If the repository exists, but is empty, returns ErrEmptyRemoteRepository.
+	AdvertisedReferencesContext(context.Context) (*packp.AdvRefs, error)
 	io.Closer
 }
 
diff --git a/plumbing/transport/http/common.go b/plumbing/transport/http/common.go
index aeedc5b..d57c0fe 100644
--- a/plumbing/transport/http/common.go
+++ b/plumbing/transport/http/common.go
@@ -3,6 +3,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net"
 	"net/http"
@@ -32,7 +33,7 @@ func applyHeadersToRequest(req *http.Request, content *bytes.Buffer, host string
 
 const infoRefsPath = "/info/refs"
 
-func advertisedReferences(s *session, serviceName string) (ref *packp.AdvRefs, err error) {
+func advertisedReferences(ctx context.Context, s *session, serviceName string) (ref *packp.AdvRefs, err error) {
 	url := fmt.Sprintf(
 		"%s%s?service=%s",
 		s.endpoint.String(), infoRefsPath, serviceName,
@@ -45,7 +46,7 @@ func advertisedReferences(s *session, serviceName string) (ref *packp.AdvRefs, e
 
 	s.ApplyAuthToRequest(req)
 	applyHeadersToRequest(req, nil, s.endpoint.Host, serviceName)
-	res, err := s.client.Do(req)
+	res, err := s.client.Do(req.WithContext(ctx))
 	if err != nil {
 		return nil, err
 	}
diff --git a/plumbing/transport/http/receive_pack.go b/plumbing/transport/http/receive_pack.go
index 433dfcf..4d14ff2 100644
--- a/plumbing/transport/http/receive_pack.go
+++ b/plumbing/transport/http/receive_pack.go
@@ -25,7 +25,11 @@ func newReceivePackSession(c *http.Client, ep *transport.Endpoint, auth transpor
 }
 
 func (s *rpSession) AdvertisedReferences() (*packp.AdvRefs, error) {
-	return advertisedReferences(s.session, transport.ReceivePackServiceName)
+	return advertisedReferences(context.TODO(), s.session, transport.ReceivePackServiceName)
+}
+
+func (s *rpSession) AdvertisedReferencesContext(ctx context.Context) (*packp.AdvRefs, error) {
+	return advertisedReferences(ctx, s.session, transport.ReceivePackServiceName)
 }
 
 func (s *rpSession) ReceivePack(ctx context.Context, req *packp.ReferenceUpdateRequest) (
diff --git a/plumbing/transport/http/upload_pack.go b/plumbing/transport/http/upload_pack.go
index db37089..e735b3d 100644
--- a/plumbing/transport/http/upload_pack.go
+++ b/plumbing/transport/http/upload_pack.go
@@ -25,7 +25,11 @@ func newUploadPackSession(c *http.Client, ep *transport.Endpoint, auth transport
 }
 
 func (s *upSession) AdvertisedReferences() (*packp.AdvRefs, error) {
-	return advertisedReferences(s.session, transport.UploadPackServiceName)
+	return advertisedReferences(context.TODO(), s.session, transport.UploadPackServiceName)
+}
+
+func (s *upSession) AdvertisedReferencesContext(ctx context.Context) (*packp.AdvRefs, error) {
+	return advertisedReferences(ctx, s.session, transport.UploadPackServiceName)
 }
 
 func (s *upSession) UploadPack(
diff --git a/plumbing/transport/http/upload_pack_test.go b/plumbing/transport/http/upload_pack_test.go
index 6fae443..f9710ab 100644
--- a/plumbing/transport/http/upload_pack_test.go
+++ b/plumbing/transport/http/upload_pack_test.go
@@ -1,8 +1,10 @@
 package http
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
+	"net/url"
 	"os"
 	"path/filepath"
 
@@ -103,3 +105,32 @@ func (s *UploadPackSuite) TestAdvertisedReferencesRedirectSchema(c *C) {
 	url := session.(*upSession).endpoint.String()
 	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
 }
+
+func (s *UploadPackSuite) TestAdvertisedReferencesContext(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")
+
+	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
+	c.Assert(err, IsNil)
+
+	info, err := session.AdvertisedReferencesContext(ctx)
+	c.Assert(err, IsNil)
+	c.Assert(info, NotNil)
+
+	url := session.(*upSession).endpoint.String()
+	c.Assert(url, Equals, "https://github.com/git-fixtures/basic")
+}
+
+func (s *UploadPackSuite) TestAdvertisedReferencesContextCanceled(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	endpoint, _ := transport.NewEndpoint("http://github.com/git-fixtures/basic")
+
+	session, err := s.Client.NewUploadPackSession(endpoint, s.EmptyAuth)
+	c.Assert(err, IsNil)
+
+	info, err := session.AdvertisedReferencesContext(ctx)
+	c.Assert(err, DeepEquals, &url.Error{Op: "Get", URL: "http://github.com/git-fixtures/basic/info/refs?service=git-upload-pack", Err: context.Canceled})
+	c.Assert(info, IsNil)
+}
diff --git a/plumbing/transport/internal/common/common.go b/plumbing/transport/internal/common/common.go
index 89432e3..75405c7 100644
--- a/plumbing/transport/internal/common/common.go
+++ b/plumbing/transport/internal/common/common.go
@@ -162,14 +162,18 @@ func (c *client) listenFirstError(r io.Reader) chan string {
 	return errLine
 }
 
-// AdvertisedReferences retrieves the advertised references from the server.
 func (s *session) AdvertisedReferences() (*packp.AdvRefs, error) {
+	return s.AdvertisedReferencesContext(context.TODO())
+}
+
+// AdvertisedReferences retrieves the advertised references from the server.
+func (s *session) AdvertisedReferencesContext(ctx context.Context) (*packp.AdvRefs, error) {
 	if s.advRefs != nil {
 		return s.advRefs, nil
 	}
 
 	ar := packp.NewAdvRefs()
-	if err := ar.Decode(s.Stdout); err != nil {
+	if err := ar.Decode(s.StdoutContext(ctx)); err != nil {
 		if err := s.handleAdvRefDecodeError(err); err != nil {
 			return nil, err
 		}
@@ -237,7 +241,7 @@ func (s *session) UploadPack(ctx context.Context, req *packp.UploadPackRequest)
 		return nil, err
 	}
 
-	if _, err := s.AdvertisedReferences(); err != nil {
+	if _, err := s.AdvertisedReferencesContext(ctx); err != nil {
 		return nil, err
 	}
 
diff --git a/plumbing/transport/server/server.go b/plumbing/transport/server/server.go
index 727f902..6f89ec3 100644
--- a/plumbing/transport/server/server.go
+++ b/plumbing/transport/server/server.go
@@ -108,6 +108,10 @@ type upSession struct {
 }
 
 func (s *upSession) AdvertisedReferences() (*packp.AdvRefs, error) {
+	return s.AdvertisedReferencesContext(context.TODO())
+}
+
+func (s *upSession) AdvertisedReferencesContext(ctx context.Context) (*packp.AdvRefs, error) {
 	ar := packp.NewAdvRefs()
 
 	if err := s.setSupportedCapabilities(ar.Capabilities); err != nil {
@@ -204,6 +208,10 @@ type rpSession struct {
 }
 
 func (s *rpSession) AdvertisedReferences() (*packp.AdvRefs, error) {
+	return s.AdvertisedReferencesContext(context.TODO())
+}
+
+func (s *rpSession) AdvertisedReferencesContext(ctx context.Context) (*packp.AdvRefs, error) {
 	ar := packp.NewAdvRefs()
 
 	if err := s.setSupportedCapabilities(ar.Capabilities); err != nil {