Merge pull request #3 from fosrl/dev

Allow use of env vars, docs, and ping interval

Author: oschwartz10612

README.md

@@ -46,6 +46,22 @@ Example:
 --endpoint https://example.com
 ```
 
+You can also run it with Docker compose. For example, a service in your `docker-compose.yml` might look like this using environment vars (recommended):
+
+```yaml
+services:
+  newt:
+    image: fosrl/newt
+    container_name: newt
+    restart: unless-stopped
+    environment:
+      - PANGOLIN_ENDPOINT=https://example.com
+      - NEWT_ID=2ix2t8xk22ubpfy 
+      - NEWT_SECRET=nnisrfsdfc7prqsp9ewo1dvtvci50j5uiqotez00dgap0ii2 
+```
+
+You can also pass the CLI args to the container:
+
 ```yaml
 services:
   newt:

docker-compose.yml

@@ -0,0 +1,9 @@
+services:
+  newt:
+    image: fosrl/newt:latest
+    container_name: newt
+    restart: unless-stopped
+    environment:
+      - PANGOLIN_ENDPOINT=https://example.com
+      - NEWT_ID=2ix2t8xk22ubpfy 
+      - NEWT_SECRET=nnisrfsdfc7prqsp9ewo1dvtvci50j5uiqotez00dgap0ii2 

entrypoint.sh

@@ -1,21 +1,10 @@
 #!/bin/sh
 
-# Sample from https://github.com/traefik/traefik-library-image/blob/5070edb25b03cca6802d75d5037576c840f73fdd/v3.1/alpine/entrypoint.sh
-
 set -e
 
 # first arg is `-f` or `--some-option`
 if [ "${1#-}" != "$1" ]; then
     set -- newt "$@"
 fi
 
-# if our command is a valid newt subcommand, let's invoke it through newt instead
-# (this allows for "docker run newt version", etc)
-if newt "$1" --help >/dev/null 2>&1
-then
-    set -- newt "$@"
-else
-    echo "= '$1' is not a newt command: assuming shell execution." 1>&2
-fi
-
 exec "$@"

main.go

@@ -112,6 +112,26 @@ func ping(tnet *netstack.Net, dst string) error {
 	return nil
 }
 
+func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{}) {
+	ticker := time.NewTicker(10 * time.Second)
+	defer ticker.Stop()
+
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				err := ping(tnet, serverIP)
+				if err != nil {
+					logger.Warn("Periodic ping failed: %v", err)
+				}
+			case <-stopChan:
+				logger.Info("Stopping ping check")
+				return
+			}
+		}
+	}()
+}
+
 func pingWithRetry(tnet *netstack.Net, dst string) error {
 	const (
 		maxAttempts = 5
@@ -222,13 +242,6 @@ func resolveDomain(domain string) (string, error) {
 	return ipAddr, nil
 }
 
-func getEnvWithDefault(key, defaultValue string) string {
-	if value := os.Getenv(key); value != "" {
-		return value
-	}
-	return defaultValue
-}
-
 func main() {
 	var (
 		endpoint   string
@@ -240,12 +253,28 @@ func main() {
 		logLevel   string
 	)
 
-	// Define CLI flags with default values from environment variables
-	flag.StringVar(&endpoint, "endpoint", os.Getenv("PANGOLIN_ENDPOINT"), "Endpoint of your pangolin server")
-	flag.StringVar(&id, "id", os.Getenv("NEWT_ID"), "Newt ID")
-	flag.StringVar(&secret, "secret", os.Getenv("NEWT_SECRET"), "Newt secret")
-	flag.StringVar(&dns, "dns", getEnvWithDefault("DEFAULT_DNS", "8.8.8.8"), "DNS server to use")
-	flag.StringVar(&logLevel, "log-level", getEnvWithDefault("LOG_LEVEL", "INFO"), "Log level (DEBUG, INFO, WARN, ERROR, FATAL)")
+	// if PANGOLIN_ENDPOINT, NEWT_ID, and NEWT_SECRET are set as environment variables, they will be used as default values
+	endpoint = os.Getenv("PANGOLIN_ENDPOINT")
+	id = os.Getenv("NEWT_ID")
+	secret = os.Getenv("NEWT_SECRET")
+	dns = os.Getenv("DNS")
+	logLevel = os.Getenv("LOG_LEVEL")
+
+	if endpoint == "" {
+		flag.StringVar(&endpoint, "endpoint", "", "Endpoint of your pangolin server")
+	}
+	if id == "" {
+		flag.StringVar(&id, "id", "", "Newt ID")
+	}
+	if secret == "" {
+		flag.StringVar(&secret, "secret", "", "Newt secret")
+	}
+	if dns == "" {
+		flag.StringVar(&dns, "dns", "8.8.8.8", "DNS server to use")
+	}
+	if logLevel == "" {
+		flag.StringVar(&logLevel, "log-level", "INFO", "Log level (DEBUG, INFO, WARN, ERROR, FATAL)")
+	}
 	flag.Parse()
 
 	logger.Init()
@@ -291,6 +320,9 @@ func main() {
 		client.Close()
 	})
 
+	pingStopChan := make(chan struct{})
+	defer close(pingStopChan)
+
 	// Register handlers for different message types
 	client.RegisterHandler("newt/wg/connect", func(msg websocket.WSMessage) {
 		logger.Info("Received registration message")
@@ -365,6 +397,11 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey(
 			logger.Error("Failed to ping %s: %v", wgData.ServerIP, err)
 		}
 
+		if !connected {
+			logger.Info("Starting ping check")
+			startPingCheck(tnet, wgData.ServerIP, pingStopChan)
+		}
+
 		// Create proxy manager
 		pm = proxy.NewProxyManager(tnet)