safesocket: add ConnectContext

This adds a variant for Connect that takes in a context.Context
which allows passing through cancellation etc by the caller.

Updates tailscale/corp#18266

Signed-off-by: Maisem Ali <maisem@tailscale.com>
This commit is contained in:
Maisem Ali
2024-06-10 19:38:10 -07:00
committed by Maisem Ali
parent 3672f66c74
commit 4b6a0c42c8
9 changed files with 32 additions and 17 deletions

View File

@@ -167,7 +167,7 @@ func (s *FileSystemForRemote) buildChild(share *drive.Share) *compositedav.Child
return fmt.Sprintf("http://%s/%s/%s", hex.EncodeToString([]byte(share.Name)), secretToken, url.PathEscape(share.Name)), nil
},
Transport: &http.Transport{
Dial: func(_, shareAddr string) (net.Conn, error) {
DialContext: func(ctx context.Context, _, shareAddr string) (net.Conn, error) {
shareNameHex, _, err := net.SplitHostPort(shareAddr)
if err != nil {
return nil, fmt.Errorf("unable to parse share address %v: %w", shareAddr, err)
@@ -188,10 +188,11 @@ func (s *FileSystemForRemote) buildChild(share *drive.Share) *compositedav.Child
_, err = netip.ParseAddrPort(addr)
if err == nil {
// this is a regular network address, dial normally
return net.Dial("tcp", addr)
var std net.Dialer
return std.DialContext(ctx, "tcp", addr)
}
// assume this is a safesocket address
return safesocket.Connect(addr)
return safesocket.ConnectContext(ctx, addr)
},
},
}