1package replica
 2
 3import (
 4	"context"
 5	"fmt"
 6	"net"
 7	"net/url"
 8	"time"
 9
10	"github.com/go-git/go-git/v5/plumbing/transport"
11	"golang.org/x/net/proxy"
12)
13
14const deadlineScheme = "tug-replica-deadline"
15
16func init() {
17	proxy.RegisterDialerType(deadlineScheme, newDeadlineDialer)
18}
19
20type deadlineDialer struct {
21	total time.Duration
22	fetch *fetchLimit
23}
24
25func newDeadlineDialer(u *url.URL, _ proxy.Dialer) (proxy.Dialer, error) {
26	if id := u.Query().Get("fetch"); id != "" {
27		limit, ok := activeFetches.Load(id)
28		if !ok {
29			return nil, fmt.Errorf("replica: fetch has ended")
30		}
31		return deadlineDialer{total: limit.(*fetchLimit).total, fetch: limit.(*fetchLimit)}, nil
32	}
33	d, err := time.ParseDuration(u.Query().Get("d"))
34	if err != nil {
35		return nil, err
36	}
37	return deadlineDialer{total: d}, nil
38}
39
40func (d deadlineDialer) Dial(network, addr string) (net.Conn, error) {
41	return d.DialContext(context.Background(), network, addr)
42}
43
44func (d deadlineDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
45	if d.fetch != nil {
46		ctx = d.fetch.ctx
47	}
48	var nd net.Dialer
49	c, err := nd.DialContext(ctx, network, addr)
50	if err != nil {
51		return nil, err
52	}
53	if d.total > 0 {
54		if err := c.SetDeadline(time.Now().Add(d.total)); err != nil {
55			c.Close()
56			return nil, err
57		}
58	}
59	if d.fetch != nil {
60		return &fetchConn{Conn: c, limit: d.fetch, stop: context.AfterFunc(ctx, func() { c.Close() })}, nil
61	}
62	return c, nil
63}
64
65func proxyFor(auth transport.AuthMethod) transport.ProxyOptions {
66	a, ok := auth.(*SSH)
67	if ok && a.fetch != "" {
68		return transport.ProxyOptions{URL: deadlineScheme + "://deadline/?fetch=" + a.fetch}
69	}
70	if !ok || a.Timeout <= 0 {
71		return transport.ProxyOptions{}
72	}
73	return transport.ProxyOptions{URL: deadlineScheme + "://deadline/?d=" + a.Timeout.String()}
74}