move thumbnail to utils

Hee Sung Son committed Apr 7, 2026 at 21:26 UTC 39a346c54adbbc8248f13d9e2fb307914f166350
2 files changed +44 -51
portal/server.go
+7 -4
@@ -22,6 +22,7 @@ import (
22 "github.com/gosuda/portal-tunnel/v2/portal/keyless"
23 "github.com/gosuda/portal-tunnel/v2/portal/policy"
24 "github.com/gosuda/portal-tunnel/v2/portal/transport"
25 + "github.com/gosuda/portal-tunnel/v2/portal/utils/thumbnail"
26 "github.com/gosuda/portal-tunnel/v2/types"
27 "github.com/gosuda/portal-tunnel/v2/utils"
28 )
@@ -70,7 +71,7 @@ type Server struct {
71 cfg ServerConfig
72 trustedProxyCIDRs []*net.IPNet
73 relaySet *discovery.RelaySet
73 - thumbnails *thumbnailService
74 + thumbnails *thumbnail.Service
75 shutdownOnce sync.Once
76 }
77
@@ -152,7 +153,7 @@ func NewServer(cfg ServerConfig) (*Server, error) {
153 tcpPorts: tcpPorts,
154 identity: identity,
155 trustedProxyCIDRs: trustedProxyCIDRs,
155 - thumbnails: newThumbnailService(cfg.HeadlessShellURL),
156 + thumbnails: thumbnail.NewService(cfg.HeadlessShellURL),
157 }
158
159 if cfg.DiscoveryEnabled {
@@ -333,8 +334,10 @@ func (s *Server) LeaseSnapshots() []types.Lease {
334 if snap.Ready == 0 && since >= 3*time.Minute {
335 continue
336 }
336 - if snap.Metadata.Thumbnail == "" && s.thumbnails != nil && s.thumbnails.Has(snap.Hostname) {
337 - snap.Metadata.Thumbnail = types.PathThumbnailPrefix + snap.Hostname
337 + if snap.Metadata.Thumbnail == "" && s.thumbnails != nil {
338 + if _, _, ok := s.thumbnails.Get(snap.Hostname); ok {
339 + snap.Metadata.Thumbnail = types.PathThumbnailPrefix + snap.Hostname
340 + }
341 }
342 out = append(out, snap.Lease)
343 }
portal/utils/thumbnail/thumbnail.go renamed
+37 -47
@@ -1,4 +1,4 @@
1 -package portal
1 +package thumbnail
2
3 import (
4 "context"
@@ -17,14 +17,14 @@ import (
17 )
18
19 const (
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"
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"
28 )
29
30 type thumbEntry struct {
@@ -32,7 +32,7 @@ type thumbEntry struct {
32 fetchedAt time.Time
33 }
34
35 -type thumbnailService struct {
35 +type Service struct {
36 mu sync.RWMutex
37 cache map[string]*thumbEntry
38 pending map[string]bool
@@ -41,15 +41,15 @@ type thumbnailService struct {
41 done chan struct{}
42 }
43
44 -func newThumbnailService(headlessShellURL string) *thumbnailService {
44 +func NewService(headlessShellURL string) *Service {
45 headlessShellURL = strings.TrimSpace(headlessShellURL)
46 if headlessShellURL == "" {
47 return nil
48 }
49 - ts := &thumbnailService{
49 + ts := &Service{
50 cache: make(map[string]*thumbEntry),
51 pending: make(map[string]bool),
52 - queue: make(chan string, thumbnailQueueSize),
52 + queue: make(chan string, queueSize),
53 headlessShellURL: headlessShellURL,
54 done: make(chan struct{}),
55 }
@@ -57,7 +57,7 @@ func newThumbnailService(headlessShellURL string) *thumbnailService {
57 return ts
58 }
59
60 -func (ts *thumbnailService) worker() {
60 +func (ts *Service) worker() {
61 for hostname := range ts.queue {
62 ts.capture(hostname)
63 ts.mu.Lock()
@@ -67,7 +67,7 @@ func (ts *thumbnailService) worker() {
67 close(ts.done)
68 }
69
70 -func (ts *thumbnailService) Get(hostname string) ([]byte, string, bool) {
70 +func (ts *Service) Get(hostname string) ([]byte, string, bool) {
71 if ts == nil {
72 return nil, "", false
73 }
@@ -77,20 +77,10 @@ func (ts *thumbnailService) Get(hostname string) ([]byte, string, bool) {
77 if !ok || len(entry.data) == 0 {
78 return nil, "", false
79 }
80 - return entry.data, thumbnailContentType, true
80 + return entry.data, ContentType, true
81 }
82
83 -func (ts *thumbnailService) Has(hostname string) bool {
84 - if ts == nil {
85 - return false
86 - }
87 - ts.mu.RLock()
88 - entry, ok := ts.cache[hostname]
89 - ts.mu.RUnlock()
90 - return ok && len(entry.data) > 0
91 -}
92 -
93 -func (ts *thumbnailService) TriggerAsync(hostname string) {
83 +func (ts *Service) TriggerAsync(hostname string) {
84 if ts == nil || hostname == "" {
85 return
86 }
@@ -99,7 +89,7 @@ func (ts *thumbnailService) TriggerAsync(hostname string) {
89 defer ts.mu.Unlock()
90
91 if entry, ok := ts.cache[hostname]; ok {
102 - if len(entry.data) > 0 || time.Since(entry.fetchedAt) < thumbnailCooldown {
92 + if len(entry.data) > 0 || time.Since(entry.fetchedAt) < cooldown {
93 return
94 }
95 }
@@ -115,30 +105,30 @@ func (ts *thumbnailService) TriggerAsync(hostname string) {
105 }
106 }
107
118 -func (ts *thumbnailService) capture(hostname string) {
108 +func (ts *Service) capture(hostname string) {
109 + store := func(data []byte) {
110 + ts.mu.Lock()
111 + ts.cache[hostname] = &thumbEntry{data: data, fetchedAt: time.Now()}
112 + ts.mu.Unlock()
113 + }
114 +
115 data, err := ts.screenshot(hostname)
116 if err != nil {
117 log.Warn().Err(err).Str("hostname", hostname).Msg("thumbnail capture failed")
122 - ts.cacheResult(hostname, nil)
118 + store(nil)
119 return
120 }
125 - if len(data) > thumbnailMaxBytes {
121 + if len(data) > maxBytes {
122 log.Warn().Str("hostname", hostname).Int("size", len(data)).Msg("thumbnail too large, discarding")
127 - ts.cacheResult(hostname, nil)
123 + store(nil)
124 return
125 }
126
131 - ts.cacheResult(hostname, data)
127 + store(data)
128 log.Info().Str("hostname", hostname).Int("size", len(data)).Msg("thumbnail captured")
129 }
130
135 -func (ts *thumbnailService) cacheResult(hostname string, data []byte) {
136 - ts.mu.Lock()
137 - ts.cache[hostname] = &thumbEntry{data: data, fetchedAt: time.Now()}
138 - ts.mu.Unlock()
139 -}
140 -
141 -func (ts *thumbnailService) resolveCDPWebSocketURL() (string, error) {
131 +func (ts *Service) resolveCDPWebSocketURL() (string, error) {
132 parsed, err := url.Parse(ts.headlessShellURL)
133 if err != nil {
134 return "", fmt.Errorf("parse headless shell URL: %w", err)
@@ -182,7 +172,7 @@ func (ts *thumbnailService) resolveCDPWebSocketURL() (string, error) {
172 return wsURL.String(), nil
173 }
174
185 -func (ts *thumbnailService) screenshot(hostname string) ([]byte, error) {
175 +func (ts *Service) screenshot(hostname string) ([]byte, error) {
176 cdpURL, err := ts.resolveCDPWebSocketURL()
177 if err != nil {
178 return nil, err
@@ -206,27 +196,27 @@ func (ts *thumbnailService) screenshot(hostname string) ([]byte, error) {
196 defer page.Close()
197
198 _ = page.SetViewport(&proto.EmulationSetDeviceMetricsOverride{
209 - Width: thumbnailViewportWidth,
210 - Height: thumbnailViewportHeight,
199 + Width: viewportWidth,
200 + Height: viewportHeight,
201 })
202 _ = browser.IgnoreCertErrors(true)
203
204 if err := page.Navigate("https://" + hostname); err != nil {
205 return nil, err
206 }
217 - if err := page.Timeout(thumbnailPageTimeout).WaitLoad(); err != nil {
207 + if err := page.Timeout(pageTimeout).WaitLoad(); err != nil {
208 return nil, err
209 }
210 time.Sleep(1 * time.Second)
211
222 - quality := thumbnailJPEGQuality
212 + quality := jpegQuality
213 return page.Screenshot(false, &proto.PageCaptureScreenshot{
214 Format: proto.PageCaptureScreenshotFormatJpeg,
215 Quality: &quality,
216 })
217 }
218
229 -func (ts *thumbnailService) Remove(hostname string) {
219 +func (ts *Service) Remove(hostname string) {
220 if ts == nil {
221 return
222 }
@@ -236,7 +226,7 @@ func (ts *thumbnailService) Remove(hostname string) {
226 ts.mu.Unlock()
227 }
228
239 -func (ts *thumbnailService) Close() {
229 +func (ts *Service) Close() {
230 if ts == nil {
231 return
232 }