Wormhole/internal/api/visitor_helpers.go

package api

import (
	"fmt"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/wormhole/wormhole/internal/store"
)

type hostRequest struct {
	store.Host
	VisitorIDs []uint `json:"visitor_ids"`
}

type tunnelRequest struct {
	store.Tunnel
	VisitorIDs []uint `json:"visitor_ids"`
}

type hostView struct {
	store.Host
	VisitorIDs []uint `json:"visitor_ids"`
}

type tunnelView struct {
	store.Tunnel
	VisitorIDs []uint `json:"visitor_ids"`
}

func (s *Server) enrichHostView(h store.Host) (hostView, error) {
	ids, err := s.store.ListResourceVisitorIDs("host", h.ID)
	if err != nil {
		return hostView{}, err
	}
	return hostView{Host: h, VisitorIDs: ids}, nil
}

func (s *Server) enrichTunnelView(t store.Tunnel) (tunnelView, error) {
	ids, err := s.store.ListResourceVisitorIDs("tunnel", t.ID)
	if err != nil {
		return tunnelView{}, err
	}
	return tunnelView{Tunnel: t, VisitorIDs: ids}, nil
}

func (s *Server) saveResourceVisitors(resourceType string, resourceID uint, enabled bool, bindMode string, visitorIDs []uint) error {
	if !enabled {
		_ = s.store.SetResourceVisitors(resourceType, resourceID, nil)
		s.cache.SetResourceVisitors(resourceType, resourceID, nil)
		return nil
	}
	if bindMode == "" {
		bindMode = "all"
	}
	if bindMode == "selected" {
		if len(visitorIDs) == 0 {
			return fmt.Errorf("请至少选择一个访客用户")
		}
		if err := s.validateVisitorIDs(visitorIDs); err != nil {
			return err
		}
		if err := s.store.SetResourceVisitors(resourceType, resourceID, visitorIDs); err != nil {
			return err
		}
		s.cache.SetResourceVisitors(resourceType, resourceID, visitorIDs)
		return nil
	}
	_ = s.store.SetResourceVisitors(resourceType, resourceID, nil)
	s.cache.SetResourceVisitors(resourceType, resourceID, nil)
	return nil
}

func (s *Server) validateVisitorIDs(ids []uint) error {
	for _, id := range ids {
		user, err := s.store.GetUserByID(id)
		if err != nil {
			return fmt.Errorf("访客用户 #%d 不存在", id)
		}
		if user.Role != "visitor" {
			return fmt.Errorf("用户 %s 不是访客账号", user.Username)
		}
	}
	return nil
}

func writeVisitorSaveError(c *gin.Context, err error) {
	c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
}

type requestIPView struct {
	ID            uint      `json:"id"`
	ResourceType  string    `json:"resource_type"`
	ResourceID    uint      `json:"resource_id"`
	ResourceLabel string    `json:"resource_label"`
	IP            string    `json:"ip"`
	HitCount      int64     `json:"hit_count"`
	LastSeenAt    time.Time `json:"last_seen_at"`
}

func requestIPViews(st *store.Store, items []store.RequestIP) []requestIPView {
	out := make([]requestIPView, 0, len(items))
	for _, item := range items {
		label := visitorResourceLabel(st, item.ResourceType, item.ResourceID)
		if label == "" && item.ResourceID > 0 {
			label = fmt.Sprintf("#%d", item.ResourceID)
		}
		out = append(out, requestIPView{
			ID:            item.ID,
			ResourceType:  item.ResourceType,
			ResourceID:    item.ResourceID,
			ResourceLabel: label,
			IP:            item.IP,
			HitCount:      item.HitCount,
			LastSeenAt:    item.LastSeenAt,
		})
	}
	return out
}

func visitorResourceLabel(st *store.Store, resourceType string, resourceID uint) string {
	switch resourceType {
	case "host":
		h, err := st.GetHostByID(resourceID)
		if err != nil {
			return ""
		}
		label := h.Host
		if h.Location != "" && h.Location != "/" {
			label += h.Location
		}
		return label
	case "tunnel":
		t, err := st.GetTunnelByID(resourceID)
		if err != nil {
			return ""
		}
		if t.Name != "" {
			return t.Name + " :" + strconv.Itoa(t.ListenPort)
		}
		return ":" + strconv.Itoa(t.ListenPort)
	default:
		return ""
	}
}