refact: move thumbnail logic to frontend

Kim committed Apr 9, 2026 at 15:44 UTC 50f8d159188de1c1ddf70d2717cf6de00483e454
6 files changed +163 -118
cmd/relay-server/admin.go
+3 -1
@@ -181,10 +181,12 @@ func (f *Frontend) serveAdmin(w http.ResponseWriter, r *http.Request) {
181 if !utils.RequireMethod(w, r, http.MethodGet) {
182 return
183 }
184 + leases := f.server.AdminLeaseSnapshots()
185 + f.attachAutomaticAdminThumbnails(leases)
186 utils.WriteAPIData(w, http.StatusOK, types.AdminSnapshotResponse{
187 ApprovalMode: string(runtime.Approver().Mode()),
188 LandingPageEnabled: f.isLandingPageEnabled(),
187 - Leases: f.server.AdminLeaseSnapshots(),
189 + Leases: leases,
190 UDP: types.AdminUDPSettingsResponse{
191 Enabled: runtime.IsUDPEnabled(),
192 MaxLeases: runtime.UDPMaxLeases(),
cmd/relay-server/frontend.go
+73 -1
@@ -36,13 +36,14 @@ type Frontend struct {
36 server *portal.Server
37 auth *adminAuth
38 adminSettingsPath string
39 + thumbnails *thumbnailService
40
41 cachedPortalHTML []byte
42 cachedPortalHTMLOnce sync.Once
43 landingPageEnabled atomic.Bool
44 }
45
45 -func NewFrontend(server *portal.Server, adminSecret string, adminSettingsPath string, defaultLandingPageEnabled bool) (*Frontend, error) {
46 +func NewFrontend(server *portal.Server, adminSecret string, adminSettingsPath string, defaultLandingPageEnabled bool, headlessShellURL string) (*Frontend, error) {
47 if server == nil {
48 return nil, errors.New("frontend requires portal server")
49 }
@@ -60,6 +61,7 @@ func NewFrontend(server *portal.Server, adminSecret string, adminSettingsPath st
61 server: server,
62 auth: newAdminAuth(adminSecret),
63 adminSettingsPath: strings.TrimSpace(adminSettingsPath),
64 + thumbnails: newThumbnailService(headlessShellURL),
65 }
66 landingPageEnabled := defaultLandingPageEnabled
67 if state.LandingPageEnabled != nil {
@@ -93,6 +95,7 @@ func (f *Frontend) Handler() *http.ServeMux {
95 mux.HandleFunc(types.PathAdmin, f.serveAdmin)
96 mux.HandleFunc(types.PathAdminPrefix, f.serveAdmin)
97 mux.HandleFunc(types.PathTunnelStatus, f.serveTunnelStatus)
98 + mux.HandleFunc(types.PathThumbnailPrefix, f.serveThumbnail)
99 mux.HandleFunc(types.PathInstallShell, func(w http.ResponseWriter, r *http.Request) {
100 serveInstallScript(w, r, f.server.PortalURL(), false)
101 })
@@ -198,6 +201,7 @@ func (f *Frontend) injectServerData(htmlContent string) string {
201 var snapshots []types.Lease
202 if f.server != nil {
203 snapshots = f.server.LeaseSnapshots()
204 + f.attachAutomaticThumbnails(snapshots)
205 }
206 jsonData, err := json.Marshal(snapshots)
207 if err != nil {
@@ -229,6 +233,67 @@ func (f *Frontend) serveTunnelStatus(w http.ResponseWriter, r *http.Request) {
233 utils.WriteAPIData(w, http.StatusOK, resp)
234 }
235
236 +func (f *Frontend) serveThumbnail(w http.ResponseWriter, r *http.Request) {
237 + if !utils.RequireMethod(w, r, http.MethodGet) {
238 + return
239 + }
240 +
241 + hostname := strings.TrimPrefix(r.URL.Path, types.PathThumbnailPrefix)
242 + hostname = strings.TrimSpace(strings.ToLower(hostname))
243 + if hostname == "" || f.server == nil || f.thumbnails == nil {
244 + http.NotFound(w, r)
245 + return
246 + }
247 +
248 + snapshot, ok := f.server.LeaseSnapshotByHostname(hostname)
249 + if !ok || snapshot.Metadata.Thumbnail != "" {
250 + f.thumbnails.remove(hostname)
251 + http.NotFound(w, r)
252 + return
253 + }
254 +
255 + data, contentType, ok := f.thumbnails.get(hostname)
256 + if !ok {
257 + var err error
258 + data, contentType, err = f.thumbnails.load(hostname)
259 + if err != nil {
260 + http.NotFound(w, r)
261 + return
262 + }
263 + }
264 +
265 + w.Header().Set("Content-Type", contentType)
266 + w.Header().Set("Cache-Control", "public, max-age=300")
267 + w.WriteHeader(http.StatusOK)
268 + _, _ = w.Write(data)
269 +}
270 +
271 +func (f *Frontend) attachAutomaticThumbnails(leases []types.Lease) {
272 + if f == nil || f.thumbnails == nil {
273 + return
274 + }
275 + for i := range leases {
276 + if leases[i].Hostname == "" || leases[i].Metadata.Thumbnail != "" {
277 + continue
278 + }
279 + leases[i].Metadata.Thumbnail = types.PathThumbnailPrefix + leases[i].Hostname
280 + f.thumbnails.triggerAsync(leases[i].Hostname)
281 + }
282 +}
283 +
284 +func (f *Frontend) attachAutomaticAdminThumbnails(leases []types.AdminLease) {
285 + if f == nil || f.thumbnails == nil {
286 + return
287 + }
288 + for i := range leases {
289 + if leases[i].Hostname == "" || leases[i].Metadata.Thumbnail != "" {
290 + continue
291 + }
292 + leases[i].Metadata.Thumbnail = types.PathThumbnailPrefix + leases[i].Hostname
293 + f.thumbnails.triggerAsync(leases[i].Hostname)
294 + }
295 +}
296 +
297 func (f *Frontend) injectOGMetadata(htmlContent, title, description string) string {
298 if title == "" {
299 title = "Portal Proxy Gateway"
@@ -260,6 +325,13 @@ func (f *Frontend) setLandingPageEnabled(enabled bool) {
325 f.landingPageEnabled.Store(enabled)
326 }
327
328 +func (f *Frontend) Close() {
329 + if f == nil || f.thumbnails == nil {
330 + return
331 + }
332 + f.thumbnails.close()
333 +}
334 +
335 func getContentType(ext string) string {
336 if ct := mime.TypeByExtension(ext); ct != "" {
337 return ct
cmd/relay-server/main.go
+2 -2
@@ -178,16 +178,16 @@ func runServer(ctx context.Context, cfg relayServerConfig) error {
178 MaxPort: cfg.MaxPort,
179 UDPEnabled: cfg.UDPEnabled,
180 TCPEnabled: cfg.TCPEnabled,
181 - HeadlessShellURL: cfg.HeadlessShellURL,
181 })
182 if err != nil {
183 return fmt.Errorf("create relay server: %w", err)
184 }
185
187 - frontend, err := NewFrontend(server, cfg.AdminSecretKey, cfg.AdminSettingsPath, cfg.LandingPageEnabled)
186 + frontend, err := NewFrontend(server, cfg.AdminSecretKey, cfg.AdminSettingsPath, cfg.LandingPageEnabled, cfg.HeadlessShellURL)
187 if err != nil {
188 return fmt.Errorf("create frontend: %w", err)
189 }
190 + defer frontend.Close()
191
192 if err := server.Start(ctx, frontend.Handler()); err != nil {
193 return fmt.Errorf("start relay server: %w", err)
cmd/relay-server/thumbnail.go renamed
+85 -72
@@ -1,4 +1,4 @@
1 -package thumbnail
1 +package main
2
3 import (
4 "context"
@@ -17,119 +17,132 @@ import (
17 )
18
19 const (
20 - viewportWidth = 1280
21 - viewportHeight = 720
22 - jpegQuality = 80
23 - maxBytes = 256 << 10 // 256KB
24 - cooldown = 30 * time.Second
25 - pageTimeout = 15 * time.Second
26 - queueSize = 32
27 - ContentType = "image/jpeg"
20 + thumbnailViewportWidth = 1280
21 + thumbnailViewportHeight = 720
22 + thumbnailJPEGQuality = 80
23 + thumbnailMaxBytes = 256 << 10 // 256KB
24 + thumbnailCooldown = 30 * time.Second
25 + thumbnailPageTimeout = 15 * time.Second
26 + thumbnailQueueSize = 32
27 + thumbnailContentType = "image/jpeg"
28 )
29
30 -type thumbEntry struct {
30 +type thumbnailEntry struct {
31 data []byte
32 fetchedAt time.Time
33 }
34
35 -type Service struct {
35 +type thumbnailService struct {
36 mu sync.RWMutex
37 - cache map[string]*thumbEntry
37 + cache map[string]*thumbnailEntry
38 pending map[string]bool
39 queue chan string
40 headlessShellURL string
41 done chan struct{}
42 }
43
44 -func NewService(headlessShellURL string) *Service {
44 +func newThumbnailService(headlessShellURL string) *thumbnailService {
45 headlessShellURL = strings.TrimSpace(headlessShellURL)
46 if headlessShellURL == "" {
47 return nil
48 }
49 - ts := &Service{
50 - cache: make(map[string]*thumbEntry),
49 + service := &thumbnailService{
50 + cache: make(map[string]*thumbnailEntry),
51 pending: make(map[string]bool),
52 - queue: make(chan string, queueSize),
52 + queue: make(chan string, thumbnailQueueSize),
53 headlessShellURL: headlessShellURL,
54 done: make(chan struct{}),
55 }
56 - go ts.worker()
57 - return ts
56 + go service.worker()
57 + return service
58 }
59
60 -func (ts *Service) worker() {
61 - for hostname := range ts.queue {
62 - ts.capture(hostname)
63 - ts.mu.Lock()
64 - delete(ts.pending, hostname)
65 - ts.mu.Unlock()
60 +func (s *thumbnailService) worker() {
61 + for hostname := range s.queue {
62 + _, _ = s.captureAndStore(hostname)
63 + s.mu.Lock()
64 + delete(s.pending, hostname)
65 + s.mu.Unlock()
66 }
67 - close(ts.done)
67 + close(s.done)
68 }
69
70 -func (ts *Service) Get(hostname string) ([]byte, string, bool) {
71 - if ts == nil {
70 +func (s *thumbnailService) get(hostname string) ([]byte, string, bool) {
71 + if s == nil {
72 return nil, "", false
73 }
74 - ts.mu.RLock()
75 - entry, ok := ts.cache[hostname]
76 - ts.mu.RUnlock()
74 + s.mu.RLock()
75 + entry, ok := s.cache[hostname]
76 + s.mu.RUnlock()
77 if !ok || len(entry.data) == 0 {
78 return nil, "", false
79 }
80 - return entry.data, ContentType, true
80 + return entry.data, thumbnailContentType, true
81 }
82
83 -func (ts *Service) TriggerAsync(hostname string) {
84 - if ts == nil || hostname == "" {
83 +func (s *thumbnailService) load(hostname string) ([]byte, string, error) {
84 + if data, contentType, ok := s.get(hostname); ok {
85 + return data, contentType, nil
86 + }
87 + data, err := s.captureAndStore(hostname)
88 + if err != nil {
89 + return nil, "", err
90 + }
91 + return data, thumbnailContentType, nil
92 +}
93 +
94 +func (s *thumbnailService) triggerAsync(hostname string) {
95 + if s == nil || hostname == "" {
96 return
97 }
98
88 - ts.mu.Lock()
89 - defer ts.mu.Unlock()
99 + s.mu.Lock()
100 + defer s.mu.Unlock()
101
91 - if entry, ok := ts.cache[hostname]; ok {
92 - if len(entry.data) > 0 || time.Since(entry.fetchedAt) < cooldown {
102 + if entry, ok := s.cache[hostname]; ok {
103 + if len(entry.data) > 0 || time.Since(entry.fetchedAt) < thumbnailCooldown {
104 return
105 }
106 }
96 - if ts.pending[hostname] {
107 + if s.pending[hostname] {
108 return
109 }
110
100 - ts.pending[hostname] = true
111 + s.pending[hostname] = true
112 select {
102 - case ts.queue <- hostname:
113 + case s.queue <- hostname:
114 default:
104 - delete(ts.pending, hostname)
115 + delete(s.pending, hostname)
116 }
117 }
118
108 -func (ts *Service) capture(hostname string) {
119 +func (s *thumbnailService) captureAndStore(hostname string) ([]byte, error) {
120 store := func(data []byte) {
110 - ts.mu.Lock()
111 - ts.cache[hostname] = &thumbEntry{data: data, fetchedAt: time.Now()}
112 - ts.mu.Unlock()
121 + s.mu.Lock()
122 + s.cache[hostname] = &thumbnailEntry{data: data, fetchedAt: time.Now()}
123 + s.mu.Unlock()
124 }
125
115 - data, err := ts.screenshot(hostname)
126 + data, err := s.screenshot(hostname)
127 if err != nil {
128 log.Warn().Err(err).Str("hostname", hostname).Msg("thumbnail capture failed")
129 store(nil)
119 - return
130 + return nil, err
131 }
121 - if len(data) > maxBytes {
122 - log.Warn().Str("hostname", hostname).Int("size", len(data)).Msg("thumbnail too large, discarding")
132 + if len(data) > thumbnailMaxBytes {
133 + err = fmt.Errorf("thumbnail too large: %d bytes", len(data))
134 + log.Warn().Err(err).Str("hostname", hostname).Int("size", len(data)).Msg("thumbnail capture failed")
135 store(nil)
124 - return
136 + return nil, err
137 }
138
139 store(data)
140 log.Info().Str("hostname", hostname).Int("size", len(data)).Msg("thumbnail captured")
141 + return data, nil
142 }
143
131 -func (ts *Service) resolveCDPWebSocketURL() (string, error) {
132 - parsed, err := url.Parse(ts.headlessShellURL)
144 +func (s *thumbnailService) resolveCDPWebSocketURL() (string, error) {
145 + parsed, err := url.Parse(s.headlessShellURL)
146 if err != nil {
147 return "", fmt.Errorf("parse headless shell URL: %w", err)
148 }
@@ -172,8 +185,8 @@ func (ts *Service) resolveCDPWebSocketURL() (string, error) {
185 return wsURL.String(), nil
186 }
187
175 -func (ts *Service) screenshot(hostname string) ([]byte, error) {
176 - cdpURL, err := ts.resolveCDPWebSocketURL()
188 +func (s *thumbnailService) screenshot(hostname string) ([]byte, error) {
189 + cdpURL, err := s.resolveCDPWebSocketURL()
190 if err != nil {
191 return nil, err
192 }
@@ -196,44 +209,44 @@ func (ts *Service) screenshot(hostname string) ([]byte, error) {
209 defer page.Close()
210
211 _ = page.SetViewport(&proto.EmulationSetDeviceMetricsOverride{
199 - Width: viewportWidth,
200 - Height: viewportHeight,
212 + Width: thumbnailViewportWidth,
213 + Height: thumbnailViewportHeight,
214 })
215 _ = browser.IgnoreCertErrors(true)
216
217 if err := page.Navigate("https://" + hostname); err != nil {
218 return nil, err
219 }
207 - if err := page.Timeout(pageTimeout).WaitLoad(); err != nil {
220 + if err := page.Timeout(thumbnailPageTimeout).WaitLoad(); err != nil {
221 return nil, err
222 }
223 time.Sleep(1 * time.Second)
224
212 - quality := jpegQuality
225 + quality := thumbnailJPEGQuality
226 return page.Screenshot(false, &proto.PageCaptureScreenshot{
227 Format: proto.PageCaptureScreenshotFormatJpeg,
228 Quality: &quality,
229 })
230 }
231
219 -func (ts *Service) Remove(hostname string) {
220 - if ts == nil {
232 +func (s *thumbnailService) remove(hostname string) {
233 + if s == nil {
234 return
235 }
223 - ts.mu.Lock()
224 - delete(ts.cache, hostname)
225 - delete(ts.pending, hostname)
226 - ts.mu.Unlock()
236 + s.mu.Lock()
237 + delete(s.cache, hostname)
238 + delete(s.pending, hostname)
239 + s.mu.Unlock()
240 }
241
229 -func (ts *Service) Close() {
230 - if ts == nil {
242 +func (s *thumbnailService) close() {
243 + if s == nil {
244 return
245 }
233 - close(ts.queue)
234 - <-ts.done
235 - ts.mu.Lock()
236 - ts.cache = make(map[string]*thumbEntry)
237 - ts.pending = make(map[string]bool)
238 - ts.mu.Unlock()
246 + close(s.queue)
247 + <-s.done
248 + s.mu.Lock()
249 + s.cache = make(map[string]*thumbnailEntry)
250 + s.pending = make(map[string]bool)
251 + s.mu.Unlock()
252 }
portal/api_server.go
-32
@@ -127,38 +127,11 @@ func (s *Server) apiHandler(base *http.ServeMux, keylessSignerHandler http.Handl
127 }
128 keylessSignerHandler.ServeHTTP(w, r)
129 default:
130 - if strings.HasPrefix(r.URL.Path, types.PathThumbnailPrefix) {
131 - s.serveThumbnail(w, r)
132 - return
133 - }
130 base.ServeHTTP(w, r)
131 }
132 })
133 }
134
139 -func (s *Server) serveThumbnail(w http.ResponseWriter, r *http.Request) {
140 - if !utils.RequireMethod(w, r, http.MethodGet) {
141 - return
142 - }
143 - hostname := strings.TrimPrefix(r.URL.Path, types.PathThumbnailPrefix)
144 - hostname = strings.TrimSpace(strings.ToLower(hostname))
145 - if hostname == "" {
146 - http.NotFound(w, r)
147 - return
148 - }
149 -
150 - data, contentType, ok := s.thumbnails.Get(hostname)
151 - if !ok {
152 - http.NotFound(w, r)
153 - return
154 - }
155 -
156 - w.Header().Set("Content-Type", contentType)
157 - w.Header().Set("Cache-Control", "public, max-age=300")
158 - w.WriteHeader(http.StatusOK)
159 - _, _ = w.Write(data)
160 -}
161 -
135 func (s *Server) handleRoot(w http.ResponseWriter, _ *http.Request) {
136 utils.WriteAPIData(w, http.StatusOK, map[string]any{
137 "service": "portal-relay",
@@ -660,11 +633,6 @@ func (s *Server) registerLease(req types.RegisterChallengeRequest, clientIP, rep
633 return types.RegisterResponse{}, err
634 }
635
663 - // Trigger thumbnail generation for apps that don't provide one
664 - if record.Metadata.Thumbnail == "" && s.thumbnails != nil {
665 - s.thumbnails.TriggerAsync(hostname)
666 - }
667 -
636 resp := types.RegisterResponse{
637 Identity: record.Copy(),
638 Hostname: hostname,
portal/server.go
-10
@@ -25,7 +25,6 @@ import (
25 "github.com/gosuda/portal-tunnel/v2/portal/wireguard"
26 "github.com/gosuda/portal-tunnel/v2/types"
27 "github.com/gosuda/portal-tunnel/v2/utils"
28 - "github.com/gosuda/portal-tunnel/v2/utils/thumbnail"
28 )
29
30 const (
@@ -59,7 +58,6 @@ type ServerConfig struct {
58 MaxPort int
59 UDPEnabled bool
60 TCPEnabled bool
62 - HeadlessShellURL string
61 }
62
63 type Server struct {
@@ -80,7 +78,6 @@ type Server struct {
78 cfg ServerConfig
79 trustedProxyCIDRs []*net.IPNet
80 relaySet *discovery.RelaySet
83 - thumbnails *thumbnail.Service
81 shutdownOnce sync.Once
82 }
83
@@ -210,7 +207,6 @@ func NewServer(cfg ServerConfig) (*Server, error) {
207 loadMgr: policy.NewLoadManager(),
208 identity: identity,
209 trustedProxyCIDRs: trustedProxyCIDRs,
213 - thumbnails: thumbnail.NewService(cfg.HeadlessShellURL),
210 }
211 if cfg.DiscoveryEnabled {
212 set := discovery.NewRelaySet()
@@ -388,9 +384,6 @@ func (s *Server) Shutdown(ctx context.Context) error {
384 if s.acmeManager != nil {
385 s.acmeManager.Stop()
386 }
391 - if s.thumbnails != nil {
392 - s.thumbnails.Close()
393 - }
387 })
388 return shutdownErr
389 }
@@ -603,9 +596,6 @@ func (s *Server) runLeaseJanitor(ctx context.Context, interval time.Duration) er
596 Str("address", lease.Address).
597 Msg("delete expired lease ens gasless txt")
598 }
606 - if s.thumbnails != nil {
607 - s.thumbnails.Remove(lease.Hostname)
608 - }
599 lease.Close()
600 }
601 }