/
main.go
134 lines (114 loc) · 3.19 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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package main
import (
"crypto/tls"
"encoding/base64"
"log/syslog"
"net/http"
"os"
docopt "github.com/docopt/docopt-go"
"github.com/gorilla/websocket"
mozlog "github.com/mozilla-services/go-mozlogrus"
log "github.com/sirupsen/logrus"
lSyslog "github.com/sirupsen/logrus/hooks/syslog"
"github.com/taskcluster/taskcluster/v57/internal"
"github.com/taskcluster/taskcluster/v57/tools/websocktunnel/wsproxy"
)
const usage = `Websocketunnel Server
Usage: websocktunnel [-h | --help]
Environment:
URL_PREFIX (required) URL prefix (http(s)://hostname(:port)) at which
this service is publicly exposed
PORT (optional; defaults to 80 or 443) port on which to listent
TLS_CERTIFICATE (optional; no TLS if not provided) base64-encoded TLS certificate
TLS_KEY corresponding base64-encoded TLS key
TASKCLUSTER_PROXY_SECRET_A JWT secret
TASKCLUSTER_PROXY_SECRET_B alternate JWT secret
SYSLOG_ADDR address to which to send syslog output
AUDIENCE JWT 'audience' claim
Options:
-h --help Show help`
func main() {
_, _ = docopt.ParseArgs(usage, nil, "websocktunnel "+internal.Version)
urlPrefix := os.Getenv("URL_PREFIX")
if urlPrefix == "" {
panic("URL_PREFIX is required")
}
logger := log.New()
if env := os.Getenv("ENV"); env == "production" {
// add mozlog formatter
logger.Formatter = &mozlog.MozLogFormatter{
LoggerName: "websocktunnel",
}
// add syslog hook if addr is provided
syslogAddr := os.Getenv("SYSLOG_ADDR")
if syslogAddr != "" {
hook, err := lSyslog.NewSyslogHook("udp", syslogAddr, syslog.LOG_DEBUG, "websocktunnel")
if err != nil {
panic(err)
}
logger.Hooks.Add(hook)
}
}
// Load secrets
signingSecretA := os.Getenv("TASKCLUSTER_PROXY_SECRET_A")
signingSecretB := os.Getenv("TASKCLUSTER_PROXY_SECRET_B")
// Load TLS certificates
useTLS := true
tlsKeyEnc := os.Getenv("TLS_KEY")
tlsCertEnc := os.Getenv("TLS_CERTIFICATE")
tlsKey, _ := base64.StdEncoding.DecodeString(tlsKeyEnc)
tlsCert, _ := base64.StdEncoding.DecodeString(tlsCertEnc)
cert, err := tls.X509KeyPair([]byte(tlsCert), []byte(tlsKey))
if err != nil {
logger.Error(err.Error())
useTLS = false
}
//load port
port := os.Getenv("PORT")
if port == "" {
if useTLS {
port = "443"
} else {
port = "80"
}
}
// load audience value
audience := os.Getenv("AUDIENCE")
upgrader := websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// will panic if secrets are not loaded
proxy, _ := wsproxy.New(wsproxy.Config{
Logger: logger,
Upgrader: upgrader,
JWTSecretA: []byte(signingSecretA),
JWTSecretB: []byte(signingSecretB),
URLPrefix: urlPrefix,
Audience: audience,
})
server := &http.Server{Addr: ":" + port, Handler: proxy}
defer func() {
_ = server.Close()
}()
logger.WithFields(log.Fields{
"server-addr": server.Addr,
}).Info("starting server")
// create tls config and serve
if useTLS {
config := &tls.Config{
Certificates: []tls.Certificate{cert},
}
listener, err := tls.Listen("tcp", ":"+port, config)
if err != nil {
panic(err)
}
_ = server.Serve(listener)
} else {
err = server.ListenAndServe()
if err != nil {
panic(err)
}
}
}