/
main.go
155 lines (126 loc) · 3.72 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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package main
import (
"context"
"crypto/tls"
"errors"
"flag"
"fmt"
"log"
"net/http"
"os"
"os/signal"
"runtime"
"runtime/pprof"
"syscall"
"time"
"github.com/redis/go-redis/v9"
)
type Config struct {
ListenAddr string
RedisAddr string
RedisUser string
RedisPassword string
MapFile string
TLSCertFile string
TLSKeyFile string
CPUProfile bool
}
func newRedisOpts(cfg *Config) *redis.Options {
const (
redisDialTimeout = 2 * time.Second
redisReadTimeout = 2 * time.Second
redisWriteTimeout = 2 * time.Second
redisDB = 0
)
var redisPoolSize = 20 & runtime.NumCPU()
return &redis.Options{
Addr: cfg.RedisAddr,
DB: redisDB,
DialTimeout: redisDialTimeout,
ReadTimeout: redisReadTimeout,
WriteTimeout: redisWriteTimeout,
ClientName: "redis-rest-api",
PoolSize: redisPoolSize,
PoolFIFO: true,
}
}
func newRedisClient(opts *redis.Options) *redis.Client {
client := redis.NewClient(opts)
return client
}
func main() {
log.SetFlags(log.LstdFlags | log.Lshortfile)
const (
shutdownTimeout = 10 * time.Second
serverReadTimeout = 2 * time.Second
serverReadHeaderTimeout = 2 * time.Second
serverIdleTimeout = 2 * time.Second
)
cfg := &Config{}
flag.StringVar(&cfg.ListenAddr, "listen-addr", ":8081", "address to listen on")
flag.StringVar(&cfg.RedisAddr, "redis-addr", "localhost:6379", "address of redis server")
flag.StringVar(&cfg.RedisUser, "redis-user", "default", "redis user to AUTH as")
flag.StringVar(&cfg.RedisPassword, "redis-password", "", "redis user password to AUTH with")
flag.StringVar(&cfg.MapFile, "map-file", "redis-users.json", "filepath containing user map")
flag.StringVar(&cfg.TLSCertFile, "tls-cert", "test-cert.pem", "TLS certificate file")
flag.StringVar(&cfg.TLSKeyFile, "tls-key", "test-key.pem", "TLS key file")
flag.BoolVar(&cfg.CPUProfile, "profile", false, "Create a CPU profile")
showVersion := flag.Bool("version", false, "print version and exit")
flag.Parse()
if *showVersion {
fmt.Printf("Version: %s\nBuild: %s\n", Version, Build)
os.Exit(0)
}
if cfg.CPUProfile {
f, err := os.Create("cpu.prof")
if err != nil {
log.Println("could not create CPU profile: ", err)
return
}
defer f.Close()
if err := pprof.StartCPUProfile(f); err != nil {
log.Println("could not start CPU profile: ", err)
return
}
defer pprof.StopCPUProfile()
}
ctx := context.Background()
um, err := loadUsers(cfg.MapFile)
if err != nil {
log.Println(err)
return
}
ropts := newRedisOpts(cfg)
rc := newRedisClient(ropts)
const maxBytes int64 = 1048576 // 1 MiB
mux := http.NewServeMux()
mux.HandleFunc("/", rootHandler(ctx, rc, um, authenticate))
mux.HandleFunc("/pipeline", pipelineHandler(ctx, rc, um, authenticate))
mux.HandleFunc("/multi-exec", txHandler(ctx, rc, um, authenticate))
wrappedMux := http.MaxBytesHandler(mux, maxBytes)
server := &http.Server{
Addr: cfg.ListenAddr,
Handler: wrappedMux,
ReadTimeout: serverReadTimeout,
ReadHeaderTimeout: serverReadHeaderTimeout,
IdleTimeout: serverIdleTimeout,
TLSConfig: &tls.Config{
MinVersion: tls.VersionTLS12,
},
}
go func() {
if err := server.ListenAndServeTLS(cfg.TLSCertFile, cfg.TLSKeyFile); !errors.Is(err, http.ErrServerClosed) {
log.Printf("HTTP server error: %v", err)
return
}
log.Println("Stopped serving new connections.")
}()
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
<-sigChan
shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), shutdownTimeout)
defer shutdownRelease()
if err := server.Shutdown(shutdownCtx); err != nil {
log.Printf("HTTP shutdown error: %v", err)
}
}