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
}