/
main.go
120 lines (97 loc) · 2.81 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package main
import (
"context"
"log"
"net"
"os"
"os/signal"
"syscall"
"github.com/alecthomas/kong"
"golang.org/x/sync/errgroup"
)
const desc = `Launch the git-ssh-uplift proxy.
SYNOPSIS
git-ssh-uplift [flags] [prog [arg ...]]
DESCRIPTION
For rationale, see https://github.com/saj/git-ssh-uplift
The uplift proxy may be invoked with or without positional arguments.
Positional arguments denote a command name (and its arguments) to be forked as a
child of the uplift proxy. UPLIFT_PORT is set in the environment of the child.
The uplift proxy automatically shuts down after its child terminates. Without
positional arguments, the uplift proxy will continue serving until it is
explicitly interrupted (with SIGINT or SIGTERM).
`
type Cmd struct {
Args []string `kong:"arg,optional"`
Bind BindAddress `kong:"default=':',placeholder='[host]:[port]',help='Bind the uplift proxy to a particular local address and/or TCP port. host defaults to all addresses on Linux, or 127.0.0.1 on other operating systems. port is randomly chosen if omitted.'"`
ConnsMax uint `kong:"default='10',help='Limit the maximum number of concurrent connections to the proxy (if value is positive).'"`
}
type BindAddress net.TCPAddr
var defaultBindAddressString = defaultBindAddress.String()
func (f *BindAddress) UnmarshalText(text []byte) error {
s := string(text)
// Defer to explicit user preference -- even if 0.0.0.0 or [::].
host, port, err := net.SplitHostPort(s)
if err != nil {
return err
}
if host == "" {
host = defaultBindAddressString
}
addr, err := net.ResolveTCPAddr("tcp", host+":"+port)
if err != nil {
return err
}
*f = *(*BindAddress)(addr)
return nil
}
type CmdContext struct {
Canceler context.Context
}
func (cmd *Cmd) Run(cmdcon CmdContext) error {
ctx, cancel := context.WithCancel(cmdcon.Canceler)
defer cancel()
lis, err := net.ListenTCP("tcp", (*net.TCPAddr)(&cmd.Bind))
if err != nil {
return err
}
port := lis.Addr().(*net.TCPAddr).Port
eg, ectx := errgroup.WithContext(ctx)
eg.Go(func() error { return serve(ectx, lis, newLimiter(cmd.ConnsMax)) })
if len(cmd.Args) > 0 {
eg.Go(func() error {
args := procArgs(cmd.Args)
err := args.ExecUplifted(ectx, port)
cancel()
return err
})
} else {
log.Printf("listening on port %d - ^C to exit", port)
}
return eg.Wait()
}
func init() {
log.SetFlags(0)
log.SetPrefix("git-ssh-uplift: ")
}
func main() {
os.Exit(run())
}
func run() int {
signals := make(chan os.Signal, 1)
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go func() {
<-signals
cancel()
}()
var cli Cmd
kong.Parse(&cli, kong.Description(desc))
cmdcon := CmdContext{Canceler: ctx}
if err := cli.Run(cmdcon); err != nil {
log.Print(err)
return 1
}
return 0
}