1  package web
  2  
  3  import (
  4  	"context"
  5  	"fmt"
  6  	"log"
  7  	"net/http"
  8  	"net/http/httputil"
  9  	"net/url"
 10  
 11  	dogeboxd "github.com/dogeorg/dogeboxd/pkg"
 12  	"github.com/dogeorg/dogeboxd/pkg/conductor"
 13  )
 14  
 15  func NewAdminRouter(config dogeboxd.ServerConfig, pm dogeboxd.PupManager) conductor.Service {
 16  	return AdminRouter{
 17  		config: config,
 18  		pm:     pm,
 19  		prx:    map[string]*adminProxy{},
 20  	}
 21  }
 22  
 23  type AdminRouter struct {
 24  	config dogeboxd.ServerConfig
 25  	pm     dogeboxd.PupManager
 26  	prx    map[string]*adminProxy
 27  }
 28  
 29  func (t *AdminRouter) updateProxies() {
 30  	// find pups that have admin ports
 31  	visited := map[string]bool{}
 32  	for pupid, pup := range t.pm.GetStateMap() {
 33  		for _, ui := range pup.WebUIs {
 34  			id := fmt.Sprintf("%s:%s", pupid, ui.Port)
 35  			visited[id] = true
 36  			_, exists := t.prx[id]
 37  			if !exists {
 38  				t.prx[id] = &adminProxy{
 39  					bindPort: ui.Port,
 40  					destHost: pup.IP,
 41  					destPort: ui.Internal,
 42  				}
 43  				t.prx[id].Start()
 44  			}
 45  		}
 46  	}
 47  	// close any that no longer exist
 48  	for id := range t.prx {
 49  		_, exists := visited[id]
 50  		if !exists {
 51  			t.prx[id].Stop()
 52  			delete(t.prx, id)
 53  		}
 54  	}
 55  }
 56  
 57  func (t AdminRouter) Run(started, stopped chan bool, stop chan context.Context) error {
 58  	t.updateProxies()
 59  	go func() {
 60  		go func() {
 61  		mainloop:
 62  			for {
 63  				select {
 64  				case <-stop:
 65  					break mainloop
 66  				case p, ok := <-t.pm.GetUpdateChannel():
 67  					if !ok {
 68  						break mainloop
 69  					}
 70  					if p.Event == dogeboxd.PUP_ADOPTED {
 71  						// New pup adopted, update proxies
 72  						t.updateProxies()
 73  					}
 74  				}
 75  			}
 76  		}()
 77  
 78  		started <- true
 79  		<-stop
 80  		// if we're stopping, shut down all the admin proxies
 81  		for _, prx := range t.prx {
 82  			prx.Stop()
 83  		}
 84  		stopped <- true
 85  	}()
 86  	return nil
 87  }
 88  
 89  // adminProxy does the actual proxying, AdminRouter manages these proxies.
 90  type adminProxy struct {
 91  	bindPort int
 92  	destHost string
 93  	destPort int
 94  	stop     context.CancelFunc
 95  }
 96  
 97  func (t *adminProxy) Start() {
 98  	fmt.Printf("Starting admin proxy %d -> %d\n", t.bindPort, t.destPort)
 99  	ctx, cancel := context.WithCancel(context.Background())
100  	t.stop = cancel
101  
102  	target := fmt.Sprintf("http://%s:%d", t.destHost, t.destPort)
103  	proxyURL, err := url.Parse(target)
104  	if err != nil {
105  		log.Fatalf("Failed to parse URL: %v", err)
106  	}
107  
108  	httpServer := &http.Server{
109  		Addr:    fmt.Sprintf("0.0.0.0:%d", t.bindPort),
110  		Handler: httputil.NewSingleHostReverseProxy(proxyURL),
111  	}
112  
113  	// Custom Director to set Cache-Control header
114  	proxy := httpServer.Handler.(*httputil.ReverseProxy)
115  	proxy.Director = func(req *http.Request) {
116  		req.URL.Scheme = proxyURL.Scheme
117  		req.URL.Host = proxyURL.Host
118  		if _, ok := req.Header["User-Agent"]; !ok {
119  			// explicitly disable User-Agent so it's not set automatically.
120  			req.Header.Set("User-Agent", "")
121  		}
122  
123  		req.Header.Set("Cache-Control", "private, max-age=10")
124  	}
125  
126  	// handle stopping
127  	go func() {
128  		<-ctx.Done()
129  		if err := httpServer.Shutdown(ctx); err != nil {
130  			log.Fatalf("Could not gracefully shutdown admin proxy: %v", err)
131  		}
132  	}()
133  
134  	// Start the server
135  	go func() {
136  		if err := httpServer.ListenAndServe(); err != http.ErrServerClosed {
137  			log.Fatalf("HTTP server ListenAndServe: %v", err)
138  		}
139  	}()
140  }
141  
142  func (t *adminProxy) Stop() {
143  	t.stop()
144  }