cloudflare/cloudflared
Publicmirrored from https://github.com/cloudflare/cloudflaredAvailable
cmd/sqlgateway/sqlgateway.go
148lines · modeblame
4e1df1a2Adam Chalmers6 years ago | 1 | package sqlgateway |
| 2 | | |
| 3 | import ( | |
| 4 | "database/sql" | |
| 5 | "encoding/json" | |
| 6 | "fmt" | |
| 7 | "math/rand" | |
| 8 | "net/http" | |
| 9 | "strings" | |
| 10 | "time" | |
| 11 | | |
| 12 | _ "github.com/lib/pq" | |
| 13 | cli "gopkg.in/urfave/cli.v2" | |
| 14 | | |
| 15 | "github.com/elgs/gosqljson" | |
| 16 | | |
| 17 | "github.com/gorilla/mux" | |
| 18 | "github.com/sirupsen/logrus" | |
| 19 | ) | |
| 20 | | |
| 21 | type Message struct { | |
| 22 | Connection Connection `json:"connection"` | |
| 23 | Command string `json:"command"` | |
| 24 | Params []interface{} `json:"params"` | |
| 25 | } | |
| 26 | | |
| 27 | type Connection struct { | |
| 28 | SSLMode string `json:"sslmode"` | |
| 29 | Token string `json:"token"` | |
| 30 | } | |
| 31 | | |
| 32 | type Response struct { | |
| 33 | Columns []string `json:"columns"` | |
| 34 | Rows [][]string `json:"rows"` | |
| 35 | Error string `json:"error"` | |
| 36 | } | |
| 37 | | |
| 38 | type Proxy struct { | |
| 39 | Context *cli.Context | |
| 40 | Router *mux.Router | |
| 41 | Token string | |
| 42 | User string | |
| 43 | Password string | |
| 44 | Driver string | |
| 45 | Database string | |
| 46 | Logger *logrus.Logger | |
| 47 | } | |
| 48 | | |
| 49 | func StartProxy(c *cli.Context, logger *logrus.Logger, password string) error { | |
| 50 | proxy := NewProxy(c, logger, password) | |
| 51 | | |
| 52 | logger.Infof("Starting SQL Gateway Proxy on port %s", strings.Split(c.String("url"), ":")[1]) | |
| 53 | | |
| 54 | err := http.ListenAndServe(":"+strings.Split(c.String("url"), ":")[1], proxy.Router) | |
| 55 | if err != nil { | |
| 56 | return err | |
| 57 | } | |
| 58 | | |
| 59 | return nil | |
| 60 | } | |
| 61 | | |
| 62 | func randID(n int, c *cli.Context) string { | |
| 63 | charBytes := []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890") | |
| 64 | b := make([]byte, n) | |
| 65 | for i := range b { | |
| 66 | b[i] = charBytes[rand.Intn(len(charBytes))] | |
| 67 | } | |
| 68 | return fmt.Sprintf("%s&%s", c.String("hostname"), b) | |
| 69 | } | |
| 70 | | |
| 71 | // db://user@dbname | |
| 72 | func parseInfo(input string) (string, string, string) { | |
| 73 | p1 := strings.Split(input, "://") | |
| 74 | p2 := strings.Split(p1[1], "@") | |
| 75 | return p1[0], p2[0], p2[1] | |
| 76 | } | |
| 77 | | |
| 78 | func NewProxy(c *cli.Context, logger *logrus.Logger, pass string) *Proxy { | |
| 79 | rand.Seed(time.Now().UnixNano()) | |
| 80 | driver, user, dbname := parseInfo(c.String("address")) | |
| 81 | proxy := Proxy{ | |
| 82 | Context: c, | |
| 83 | Router: mux.NewRouter(), | |
| 84 | Token: randID(64, c), | |
| 85 | Logger: logger, | |
| 86 | User: user, | |
| 87 | Password: pass, | |
| 88 | Database: dbname, | |
| 89 | Driver: driver, | |
| 90 | } | |
| 91 | | |
| 92 | logger.Info(fmt.Sprintf(` | |
| 93 | | |
| 94 | -------------------- | |
| 95 | SQL Gateway Proxy | |
| 96 | Token: %s | |
| 97 | -------------------- | |
| 98 | | |
| 99 | `, proxy.Token)) | |
| 100 | | |
| 101 | proxy.Router.HandleFunc("/", proxy.proxyRequest).Methods("POST") | |
| 102 | return &proxy | |
| 103 | } | |
| 104 | | |
| 105 | func (proxy *Proxy) proxyRequest(rw http.ResponseWriter, req *http.Request) { | |
| 106 | var message Message | |
| 107 | response := Response{} | |
| 108 | | |
| 109 | err := json.NewDecoder(req.Body).Decode(&message) | |
| 110 | if err != nil { | |
| 111 | proxy.Logger.Error(err) | |
| 112 | http.Error(rw, fmt.Sprintf("400 - %s", err.Error()), http.StatusBadRequest) | |
| 113 | return | |
| 114 | } | |
| 115 | | |
| 116 | if message.Connection.Token != proxy.Token { | |
| 117 | proxy.Logger.Error("Invalid token") | |
| 118 | http.Error(rw, "400 - Invalid token", http.StatusBadRequest) | |
| 119 | return | |
| 120 | } | |
| 121 | | |
| 122 | connStr := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=%s", proxy.User, proxy.Password, proxy.Database, message.Connection.SSLMode) | |
| 123 | | |
| 124 | db, err := sql.Open(proxy.Driver, connStr) | |
| 125 | defer db.Close() | |
| 126 | | |
| 127 | if err != nil { | |
| 128 | proxy.Logger.Error(err) | |
| 129 | http.Error(rw, fmt.Sprintf("400 - %s", err.Error()), http.StatusBadRequest) | |
| 130 | return | |
| 131 | | |
| 132 | } else { | |
| 133 | proxy.Logger.Info("Forwarding SQL: ", message.Command) | |
| 134 | rw.Header().Set("Content-Type", "application/json") | |
| 135 | | |
| 136 | headers, data, err := gosqljson.QueryDbToArray(db, "lower", message.Command, message.Params...) | |
| 137 | | |
| 138 | if err != nil { | |
| 139 | proxy.Logger.Error(err) | |
| 140 | http.Error(rw, fmt.Sprintf("400 - %s", err.Error()), http.StatusBadRequest) | |
| 141 | return | |
| 142 | | |
| 143 | } else { | |
| 144 | response = Response{headers, data, ""} | |
| 145 | } | |
| 146 | } | |
| 147 | json.NewEncoder(rw).Encode(response) | |
| 148 | } |