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 "" } }