From ebb45b55c317da65b5884fc785a9943e2ff2f15c Mon Sep 17 00:00:00 2001
From: Christopher LaPointe <cmlapointe11@gmail.com>
Date: Fri, 26 May 2023 19:52:17 -0400
Subject: [PATCH] Handle signals, code cleanup, easier labels

---
 .dockerignore            |  2 +-
 .gitignore               |  6 ++++-
 assets.go                |  6 ++---
 assets/splash.html       |  4 +--
 docker-compose.yml       | 57 ----------------------------------------
 main.go                  | 32 +++++++++++++++++-----
 pkg/service/container.go | 46 ++++++++++++++++----------------
 pkg/service/labels.go    | 53 +++++++++++++++++++++++++++++++++++++
 pkg/service/service.go   | 28 +++++++++-----------
 pkg/service/util.go      | 13 +--------
 10 files changed, 123 insertions(+), 124 deletions(-)
 delete mode 100644 docker-compose.yml
 create mode 100644 pkg/service/labels.go

diff --git a/.dockerignore b/.dockerignore
index ee647f5..59c7e3d 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,5 +1,5 @@
 README.md
-docker-compose.yml
+*compose.yml
 Dockerfile
 .gitignore
 traefik-lazyload
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index de4cb2e..75dbf77 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,5 @@
-traefik-lazyload
\ No newline at end of file
+# Binary
+traefik-lazyload
+
+# Various testing/compose files
+*compose.yml
\ No newline at end of file
diff --git a/assets.go b/assets.go
index 96ea09f..6ee5713 100644
--- a/assets.go
+++ b/assets.go
@@ -14,10 +14,8 @@ var httpAssets embed.FS
 const httpAssetPrefix = "/__llassets/"
 
 type SplashModel struct {
-	Name        string
-	CID         string
-	WaitForCode int
-	WaitForPath string
+	*service.ContainerState
+	Name string
 }
 
 var splashTemplate = template.Must(template.ParseFS(httpAssets, path.Join("assets", config.Model.Splash)))
diff --git a/assets/splash.html b/assets/splash.html
index f10455a..69d6d5d 100644
--- a/assets/splash.html
+++ b/assets/splash.html
@@ -23,13 +23,13 @@
         </div>
         <div class="message">
             <h2>Starting {{.Name}}</h2>
-            <h3>{{.CID}}</h3>
+            <h3>{{.Name}}</h3>
         </div>
     </div>
     <script>
         async function testForOk(url) {
             const response = await fetch(url, {
-                method: "HEAD",
+                method: "{{.WaitForMethod}}",
             });
             console.log(`Got ${response.status}`);
             return response.status === {{.WaitForCode}};
diff --git a/docker-compose.yml b/docker-compose.yml
deleted file mode 100644
index 858c72a..0000000
--- a/docker-compose.yml
+++ /dev/null
@@ -1,57 +0,0 @@
-version: '3.5'
-
-services:
-  # reverse-proxy:
-  #   image: traefik:v2.4
-  #   command:
-  #     - --api.insecure
-  #     - --providers.docker
-  #     - --providers.docker.defaultRule=Host(`{{.Name}}.d0.lan`)
-  #     - --entryPoints.web.address=:85
-  #     - --entryPoints.web.forwardedHeaders.insecure
-  #     - --providers.docker.exposedByDefault=false
-  #     - --providers.docker.constraints=Label(`my.zone`,`zone1`)
-  #   restart: always
-  #   ports:
-  #     - "85:85"     # The HTTP port
-  #     - "8086:8080" # The Web UI (enabled by --api)
-  #   volumes:
-  #     - /var/run/docker.sock:/var/run/docker.sock # So that Traefik can listen to the Docker events
-
-  monitor:
-    build: .
-    labels:
-      - traefik.enable=true
-      - "traefik.http.routers.lazyload.priority=-100"
-      - "traefik.http.routers.lazyload.rule=Host(`whoami.d.lan`, `whoami2.d.lan`, `pdf.d.lan`, `noexist.d.lan`, `lazyloader.d.lan`)"
-    environment:
-      TLL_STOPATBOOT: true
-      TLL_STATUSHOST: lazyloader.d.lan
-    networks:
-      - traefik-bridge
-    volumes:
-      - /var/run/docker.sock:/var/run/docker.sock
-
-  whoami:
-    image: containous/whoami
-    networks:
-      - traefik-bridge
-    labels:
-      - traefik.enable=true
-      - "traefik.http.routers.lazywhoami.rule=Host(`whoami.d.lan`)"
-      - lazyloader=true
-
-  whoami2:
-    image: containous/whoami
-    networks:
-      - traefik-bridge
-    labels:
-      - traefik.enable=true
-      - "traefik.http.routers.lazywhoami2.rule=Host(`whoami2.d.lan`)"
-      - lazyloader=true
-      - lazyloader.stopdelay=1m
-
-networks:
-  traefik-bridge:
-    external: true
-    name: traefik-bridge
diff --git a/main.go b/main.go
index 5728862..9805808 100644
--- a/main.go
+++ b/main.go
@@ -1,11 +1,14 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"io/fs"
 	"net/http"
+	"os"
+	"os/signal"
 	"runtime"
 	"traefik-lazyload/pkg/config"
 	"traefik-lazyload/pkg/service"
@@ -44,14 +47,31 @@ func main() {
 
 	// Set up http server
 	subFs, _ := fs.Sub(httpAssets, "assets")
-	http.Handle(httpAssetPrefix, http.StripPrefix(httpAssetPrefix, http.FileServer(http.FS(subFs))))
-	http.HandleFunc("/", ContainerHandler)
+	router := http.NewServeMux()
+	router.Handle(httpAssetPrefix, http.StripPrefix(httpAssetPrefix, http.FileServer(http.FS(subFs))))
+	router.HandleFunc("/", ContainerHandler)
+
+	srv := &http.Server{
+		Addr:    config.Model.Listen,
+		Handler: router,
+	}
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt)
+
+	go func() {
+		<-sigChan
+		logrus.Info("Shutting down...")
+		srv.Shutdown(context.Background())
+	}()
 
 	logrus.Infof("Listening on %s...", config.Model.Listen)
 	if config.Model.StatusHost != "" {
 		logrus.Infof("Status host set to %s", config.Model.StatusHost)
 	}
-	http.ListenAndServe(config.Model.Listen, nil)
+	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		logrus.Fatal(err)
+	}
 }
 
 func ContainerHandler(w http.ResponseWriter, r *http.Request) {
@@ -77,10 +97,8 @@ func ContainerHandler(w http.ResponseWriter, r *http.Request) {
 	} else {
 		w.WriteHeader(http.StatusAccepted)
 		renderErr := splashTemplate.Execute(w, SplashModel{
-			Name:        host,
-			CID:         sOpts.ContainerName,
-			WaitForCode: sOpts.WaitForCode,
-			WaitForPath: sOpts.WaitForPath,
+			Name:           host,
+			ContainerState: sOpts,
 		})
 		if renderErr != nil {
 			logrus.Error(renderErr)
diff --git a/pkg/service/container.go b/pkg/service/container.go
index 449b893..b77263b 100644
--- a/pkg/service/container.go
+++ b/pkg/service/container.go
@@ -1,19 +1,19 @@
 package service
 
 import (
-	"strconv"
 	"strings"
 	"time"
 	"traefik-lazyload/pkg/config"
 
 	"github.com/docker/docker/api/types"
-	"github.com/sirupsen/logrus"
 )
 
 type containerSettings struct {
-	stopDelay   time.Duration
-	waitForCode int
-	waitForPath string
+	stopDelay     time.Duration
+	waitForCode   int
+	waitForPath   string
+	waitForMethod string
+	needs         []string
 }
 
 type ContainerState struct {
@@ -34,26 +34,11 @@ func newStateFromContainer(ct *types.Container) *ContainerState {
 }
 
 func extractContainerLabels(ct *types.Container) (target containerSettings) {
-	{ // Parse stop delay
-		stopDelay, _ := labelOrDefault(ct, "stopdelay", config.Model.StopDelay.String())
-		if dur, stopErr := time.ParseDuration(stopDelay); stopErr != nil {
-			target.stopDelay = config.Model.StopDelay
-			logrus.Warnf("Unable to parse stopdelay for %s of %s, defaulting to %s", containerShort(ct), stopDelay, target.stopDelay.String())
-		} else {
-			target.stopDelay = dur
-		}
-	}
-	{ // WaitForCode
-		codeStr, _ := labelOrDefault(ct, "waitforcode", "200")
-		if code, err := strconv.Atoi(codeStr); err != nil {
-			target.waitForCode = 200
-			logrus.Warnf("Unable to parse WaitForCode of %s, defaulting to %d", containerShort(ct), target.waitForCode)
-		} else {
-			target.waitForCode = code
-		}
-	}
-
+	target.stopDelay, _ = labelOrDefaultDuration(ct, "stopdelay", config.Model.StopDelay)
+	target.waitForCode, _ = labelOrDefaultInt(ct, "waitforcode", 200)
 	target.waitForPath, _ = labelOrDefault(ct, "waitforpath", "/")
+	target.waitForMethod, _ = labelOrDefault(ct, "waitformethod", "HEAD")
+	target.needs, _ = labelOrDefaultArr(ct, "needs")
 	return
 }
 
@@ -85,6 +70,19 @@ func (s *containerSettings) StopDelay() string {
 	return s.stopDelay.String()
 }
 
+func (s *ContainerState) WaitForCode() int {
+	return s.waitForCode
+}
+
+func (s *ContainerState) WaitForPath() string {
+	return s.waitForPath
+}
+
+func (s *ContainerState) WaitForMethod() string {
+	return s.waitForMethod
+}
+
+// Wrapper for container results that opaques and adds some methods to that data
 type ContainerWrapper struct {
 	types.Container
 }
diff --git a/pkg/service/labels.go b/pkg/service/labels.go
new file mode 100644
index 0000000..5270e39
--- /dev/null
+++ b/pkg/service/labels.go
@@ -0,0 +1,53 @@
+package service
+
+import (
+	"strconv"
+	"strings"
+	"time"
+	"traefik-lazyload/pkg/config"
+
+	"github.com/docker/docker/api/types"
+	"github.com/sirupsen/logrus"
+)
+
+func labelOrDefault(ct *types.Container, sublabel, dflt string) (string, bool) {
+	if val, ok := ct.Labels[config.SubLabel(sublabel)]; ok {
+		return val, true
+	}
+	return dflt, false
+}
+
+func labelOrDefaultArr(ct *types.Container, sublabel string) ([]string, bool) {
+	if val, ok := ct.Labels[config.SubLabel(sublabel)]; ok {
+		return strings.Split(val, ","), true
+	}
+	return []string{}, false
+}
+
+func labelOrDefaultInt(ct *types.Container, sublabel string, dflt int) (int, bool) {
+	s, ok := labelOrDefault(ct, sublabel, "")
+	if !ok {
+		return dflt, false
+	}
+
+	if val, err := strconv.Atoi(s); err != nil {
+		logrus.Warnf("Unable to parse %s on %s: %v. Using default of %d", sublabel, containerShort(ct), err, dflt)
+		return dflt, false
+	} else {
+		return val, true
+	}
+}
+
+func labelOrDefaultDuration(ct *types.Container, sublabel string, dflt time.Duration) (time.Duration, bool) {
+	s, ok := labelOrDefault(ct, sublabel, "")
+	if !ok {
+		return dflt, false
+	}
+
+	if val, err := time.ParseDuration(s); err != nil {
+		logrus.Warnf("Unable to parse %s on %s: %v. Using default of %s", sublabel, containerShort(ct), err, dflt.String())
+		return dflt, false
+	} else {
+		return val, true
+	}
+}
diff --git a/pkg/service/service.go b/pkg/service/service.go
index b55ab92..7d6f0b9 100644
--- a/pkg/service/service.go
+++ b/pkg/service/service.go
@@ -55,13 +55,7 @@ func (s *Core) Close() error {
 	return s.client.Close()
 }
 
-type StartResult struct {
-	WaitForCode   int
-	WaitForPath   string
-	ContainerName string
-}
-
-func (s *Core) StartHost(hostname string) (*StartResult, error) {
+func (s *Core) StartHost(hostname string) (*ContainerState, error) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 
@@ -75,11 +69,7 @@ func (s *Core) StartHost(hostname string) (*StartResult, error) {
 
 	if ets, exists := s.active[ct.ID]; exists {
 		logrus.Debugf("Asked to start host, but we already think it's started: %s", ets.name)
-		return &StartResult{
-			WaitForCode:   ets.waitForCode,
-			WaitForPath:   ets.waitForPath,
-			ContainerName: containerShort(ct),
-		}, nil
+		return ets, nil
 	}
 
 	go s.startContainer(ctx, ct)
@@ -88,10 +78,7 @@ func (s *Core) StartHost(hostname string) (*StartResult, error) {
 	ets := newStateFromContainer(ct)
 	s.active[ct.ID] = ets
 
-	return &StartResult{
-		WaitForCode: ets.waitForCode,
-		WaitForPath: ets.waitForPath,
-	}, nil
+	return ets, nil
 }
 
 func (s *Core) StopAll() {
@@ -229,6 +216,15 @@ func (s *Core) checkContainerForInactivity(ctx context.Context, cid string, ct *
 	return false, nil
 }
 
+func (s *Core) findContainersByDepProvider(ctx context.Context, name string) ([]types.Container, error) {
+	filters := filters.NewArgs()
+	filters.Add("label", config.SubLabel("providers")+"="+name)
+	return s.client.ContainerList(ctx, types.ContainerListOptions{
+		Filters: filters,
+		All:     true,
+	})
+}
+
 func (s *Core) findContainerByHostname(ctx context.Context, hostname string) (*types.Container, error) {
 	containers, err := s.findAllLazyloadContainers(ctx, true)
 	if err != nil {
diff --git a/pkg/service/util.go b/pkg/service/util.go
index 3025d11..b2fb2d5 100644
--- a/pkg/service/util.go
+++ b/pkg/service/util.go
@@ -3,7 +3,6 @@ package service
 import (
 	"fmt"
 	"strings"
-	"traefik-lazyload/pkg/config"
 
 	"github.com/docker/docker/api/types"
 )
@@ -16,13 +15,6 @@ func sumNetworkBytes(networks map[string]types.NetworkStats) (recv int64, send i
 	return
 }
 
-func labelOrDefault(ct *types.Container, sublabel, dflt string) (string, bool) {
-	if val, ok := ct.Labels[config.SubLabel(sublabel)]; ok {
-		return val, true
-	}
-	return dflt, false
-}
-
 func shortId(id string) string {
 	const SLEN = 8
 	if len(id) <= SLEN {
@@ -42,10 +34,7 @@ func containerShort(c *types.Container) string {
 }
 
 func trimRootPath(s string) string {
-	if strings.HasPrefix(s, "/") {
-		return s[1:]
-	}
-	return s
+	return strings.TrimPrefix(s, "/")
 }
 
 func isRunning(c *types.Container) bool {