Refator using le utils

5 files changed, 177 insertions, 0 deletions

diff --git a/util/file.go b/util/file.go
new file mode 100644
index 0000000..b6fc86f
--- /dev/null
+++ b/util/file.go
@@ -0,0 +1,19 @@
+package util
+
+import (
+        "io"
+        "io/ioutil"
+)
+
+func WriteReaderToFile(filename string, data io.Reader) error {
+        d, err := ioutil.ReadAll(data)
+        if err != nil {
+                return err
+        }
+
+        if err = ioutil.WriteFile(filename, d, 0644); err != nil {
+                return err
+        }
+
+        return nil
+}
diff --git a/util/gin.go b/util/gin.go
new file mode 100644
index 0000000..9bba6ee
--- /dev/null
+++ b/util/gin.go
@@ -0,0 +1,72 @@
+package util
+
+import (
+        "context"
+        "log"
+        "net"
+        "net/http"
+        "os"
+        "os/signal"
+
+        "github.com/gin-gonic/gin"
+)
+
+// Copied from: https://github.com/gin-gonic/gin/blob/59ab588bf597f9f41faee4f217b5659893c2e925/utils.go#L137
+func resolveAddress(addr []string) string {
+        switch len(addr) {
+        case 0:
+                if port := os.Getenv("PORT"); port != "" {
+                        log.Printf("Environment variable PORT=\"%s\"", port)
+                        return ":" + port
+                }
+                log.Printf("Environment variable PORT is undefined. Using port :8080 by default")
+                return ":8080"
+        case 1:
+                return addr[0]
+        default:
+                panic("too many parameters")
+        }
+}
+
+// Runs a gin.Engine instance in a way that can be canceled by an SIGINT
+func GinRun(e *gin.Engine, debug bool, a ...string) {
+        if debug {
+                gin.SetMode(gin.DebugMode)
+        } else {
+                gin.SetMode(gin.ReleaseMode)
+        }
+
+        srv := &http.Server{
+                Addr:    resolveAddress(a),
+                Handler: e,
+        }
+
+        idleConnsClosed := make(chan struct{})
+
+        go func() {
+                sigint := make(chan os.Signal, 1)
+                signal.Notify(sigint, os.Interrupt)
+                <-sigint
+
+                log.Println("Caught SIGINT, shutting down")
+                if err := srv.Shutdown(context.Background()); err != nil {
+                        log.Printf("HTTP server Shutdown: %v", err)
+                }
+
+                close(idleConnsClosed)
+        }()
+
+        log.Printf("Listening and serving HTTP on %s\n", srv.Addr)
+        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
+                log.Fatalf("HTTP server ListenAndServe: %v", err)
+        }
+
+        <-idleConnsClosed
+}
+
+func GetRequestIP(c *gin.Context) net.IP {
+        if xff := c.Request.Header.Get("X-Forwarded-For"); xff != "" {
+                return ParseIP(xff)
+        }
+        return ParseIP(c.Request.RemoteAddr)
+}
diff --git a/util/http.go b/util/http.go
new file mode 100644
index 0000000..31cd94d
--- /dev/null
+++ b/util/http.go
@@ -0,0 +1,26 @@
+package util
+
+import (
+        "fmt"
+        "net/http"
+        "net/url"
+)
+
+func MakeURL(r *http.Request, path string, subs ...interface{}) *url.URL {
+        scheme := "https"
+        if r.TLS == nil {
+                scheme = "http"
+        }
+
+        // Always defer to whatever the proxy told us it was doing because this
+        // could be a mullet-VIP in either direction.
+        if fwProto := r.Header.Get("X-Forwarded-Proto"); fwProto != "" {
+                scheme = fwProto
+        }
+
+        return &url.URL{
+                Scheme: scheme,
+                Host:   r.Host,
+                Path:   fmt.Sprintf(path, subs...),
+        }
+}
diff --git a/util/ip.go b/util/ip.go
new file mode 100644
index 0000000..f86d5b5
--- /dev/null
+++ b/util/ip.go
@@ -0,0 +1,24 @@
+package util
+
+import (
+        "net"
+        "regexp"
+)
+
+// Parses an IPv4 or IPv6 address with an optional port on the end. Returns
+// match groups for the addresses. The first match is the IPv6 address and the
+// second the IPv4 address.
+var ipRegexp = regexp.MustCompile(`(?:\[([0-9a-f:]+)\]|(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}))(?::\d+)?`)
+
+func ParseIP(s string) net.IP {
+        ips := ipRegexp.FindStringSubmatch(s)
+        if ips == nil {
+                return nil
+        }
+
+        if v6, v4 := ips[1], ips[2]; v6 != "" {
+                return net.ParseIP(v6)
+        } else {
+                return net.ParseIP(v4)
+        }
+}
diff --git a/util/viper.go b/util/viper.go
new file mode 100644
index 0000000..0f08a6c
--- /dev/null
+++ b/util/viper.go
@@ -0,0 +1,36 @@
+package util
+
+import (
+        "strings"
+
+        "github.com/spf13/cobra"
+        "github.com/spf13/viper"
+)
+
+type ViperWrap struct {
+        cmd *cobra.Command
+}
+
+func WrapViper(cmd *cobra.Command, prefix string) *ViperWrap {
+        viper.SetEnvPrefix(prefix)
+        viper.AutomaticEnv()
+        return &ViperWrap{cmd}
+}
+
+func (v *ViperWrap) bindViper(name, short string, defaultv interface{}, help string) {
+        vname := strings.ReplaceAll(name, "-", "_")
+        viper.BindPFlag(vname, v.cmd.Flags().Lookup(name))
+        viper.SetDefault(vname, defaultv)
+}
+
+func (v *ViperWrap) BindString(name, short, defaultv, help string) *ViperWrap {
+        v.cmd.Flags().StringP(name, short, defaultv, help)
+        v.bindViper(name, short, defaultv, help)
+        return v
+}
+
+func (v *ViperWrap) BindBool(name, short string, defaultv bool, help string) *ViperWrap {
+        v.cmd.Flags().BoolP(name, short, defaultv, help)
+        v.bindViper(name, short, defaultv, help)
+        return v
+}