Wormhole/internal/api/import_export.go

package api

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

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

type importResult struct {
	Created int      `json:"created"`
	Updated int      `json:"updated"`
	Skipped int      `json:"skipped"`
	Failed  int      `json:"failed"`
	Errors  []string `json:"errors"`
}

type tunnelImportRequest struct {
	Mode  string                   `json:"mode"` // create, upsert, skip
	Items []store.TunnelExportItem `json:"items"`
}

type hostImportRequest struct {
	Mode  string                  `json:"mode"`
	Items []store.HostExportItem  `json:"items"`
}

func (s *Server) exportTunnels(c *gin.Context) {
	claims := getClaims(c)
	clientID, _ := strconv.ParseUint(c.Query("client_id"), 10, 64)
	tunnels, err := s.store.ListTunnels(uint(clientID), claims.UserID, auth.IsAdmin(claims.Role))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	clientNames := tunnelClientNames(s.store, tunnels)
	items := make([]store.TunnelExportItem, 0, len(tunnels))
	for _, t := range tunnels {
		items = append(items, store.TunnelToExportItem(t, clientNames[t.ClientID]))
	}
	doc := store.TunnelExportDoc{
		Version:    store.ExportVersion,
		ExportedAt: time.Now().UTC(),
		Items:      items,
	}
	c.Header("Content-Disposition", `attachment; filename="wormhole-tunnels.json"`)
	c.JSON(http.StatusOK, doc)
}

func (s *Server) importTunnels(c *gin.Context) {
	claims := getClaims(c)
	var req tunnelImportRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "bad request"})
		return
	}
	if len(req.Items) == 0 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "items is empty"})
		return
	}
	mode := req.Mode
	if mode == "" {
		mode = "upsert"
	}
	result := importResult{}
	for i, item := range req.Items {
		clientID, err := store.ResolveImportClientID(s.store, item.ClientID, item.ClientName, claims.UserID, auth.IsAdmin(claims.Role))
		if err != nil {
			result.Failed++
			result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
			continue
		}
		t := tunnelFromImportItem(item, clientID)
		if err := normalizeTunnel(&t); err != nil {
			result.Failed++
			result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
			continue
		}
		existing, err := s.store.FindTunnelByPort(t.ListenPort)
		switch mode {
		case "create":
			if err == nil {
				result.Skipped++
				continue
			}
			if err := s.store.CreateTunnel(&t); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
				continue
			}
			s.cache.UpsertTunnel(&t)
			if t.Status {
				_ = s.proxy.StartTunnel(t.ID)
			}
			result.Created++
		case "skip":
			if err == nil {
				result.Skipped++
				continue
			}
			if err := s.store.CreateTunnel(&t); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
				continue
			}
			s.cache.UpsertTunnel(&t)
			if t.Status {
				_ = s.proxy.StartTunnel(t.ID)
			}
			result.Created++
		default: // upsert
			if err != nil {
				if err := s.store.CreateTunnel(&t); err != nil {
					result.Failed++
					result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
					continue
				}
				s.cache.UpsertTunnel(&t)
				if t.Status {
					_ = s.proxy.StartTunnel(t.ID)
				}
				result.Created++
				continue
			}
			wasRunning := existing.RunStatus
			if wasRunning {
				_ = s.proxy.StopTunnel(existing.ID)
			}
			t.ID = existing.ID
			t.ClientID = existing.ClientID
			t.RunStatus = false
			if err := s.store.UpdateTunnel(&t); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("tunnel", i+1, err))
				continue
			}
			s.cache.UpsertTunnel(&t)
			if t.Status {
				_ = s.proxy.StartTunnel(t.ID)
			}
			result.Updated++
		}
	}
	c.JSON(http.StatusOK, result)
}

func (s *Server) exportHosts(c *gin.Context) {
	claims := getClaims(c)
	clientID, _ := strconv.ParseUint(c.Query("client_id"), 10, 64)
	hosts, err := s.store.ListHosts(uint(clientID), claims.UserID, auth.IsAdmin(claims.Role))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	clientNames := hostClientNames(s.store, hosts)
	items := make([]store.HostExportItem, 0, len(hosts))
	for _, h := range hosts {
		items = append(items, store.HostToExportItem(h, clientNames[h.ClientID]))
	}
	doc := store.HostExportDoc{
		Version:    store.ExportVersion,
		ExportedAt: time.Now().UTC(),
		Items:      items,
	}
	c.Header("Content-Disposition", `attachment; filename="wormhole-hosts.json"`)
	c.JSON(http.StatusOK, doc)
}

func (s *Server) importHosts(c *gin.Context) {
	claims := getClaims(c)
	var req hostImportRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "bad request"})
		return
	}
	if len(req.Items) == 0 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "items is empty"})
		return
	}
	mode := req.Mode
	if mode == "" {
		mode = "upsert"
	}
	result := importResult{}
	for i, item := range req.Items {
		clientID, err := store.ResolveImportClientID(s.store, item.ClientID, item.ClientName, claims.UserID, auth.IsAdmin(claims.Role))
		if err != nil {
			result.Failed++
			result.Errors = append(result.Errors, formatImportError("host", i+1, err))
			continue
		}
		h := hostFromImportItem(item, clientID)
		if err := normalizeHost(&h); err != nil {
			result.Failed++
			result.Errors = append(result.Errors, formatImportError("host", i+1, err))
			continue
		}
		existing, err := s.store.FindHostByRouteGlobal(h.Host, h.Location)
		switch mode {
		case "create":
			if err == nil {
				result.Skipped++
				continue
			}
			if err := s.store.CreateHost(&h); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("host", i+1, err))
				continue
			}
			s.cache.UpsertHost(&h)
			result.Created++
		case "skip":
			if err == nil {
				result.Skipped++
				continue
			}
			if err := s.store.CreateHost(&h); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("host", i+1, err))
				continue
			}
			s.cache.UpsertHost(&h)
			result.Created++
		default:
			if err != nil {
				if err := s.store.CreateHost(&h); err != nil {
					result.Failed++
					result.Errors = append(result.Errors, formatImportError("host", i+1, err))
					continue
				}
				s.cache.UpsertHost(&h)
				result.Created++
				continue
			}
			h.ID = existing.ID
			h.ClientID = existing.ClientID
			if err := s.store.UpdateHost(&h); err != nil {
				result.Failed++
				result.Errors = append(result.Errors, formatImportError("host", i+1, err))
				continue
			}
			s.cache.UpsertHost(&h)
			result.Updated++
		}
	}
	c.JSON(http.StatusOK, result)
}

func tunnelFromImportItem(item store.TunnelExportItem, clientID uint) store.Tunnel {
	return store.Tunnel{
		ClientID:      clientID,
		Name:          item.Name,
		Mode:          item.Mode,
		ListenIP:      item.ListenIP,
		ListenPort:    item.ListenPort,
		TargetAddr:    item.TargetAddr,
		IPForwardMode: item.IPForwardMode,
		CustomHeaders: item.CustomHeaders,
		Status:        item.Status,
	}
}

func hostFromImportItem(item store.HostExportItem, clientID uint) store.Host {
	return store.Host{
		ClientID:      clientID,
		Host:          item.Host,
		Location:      item.Location,
		Scheme:        item.Scheme,
		TargetAddr:    item.TargetAddr,
		AutoHTTPS:     item.AutoHTTPS,
		IPForwardMode: item.IPForwardMode,
		CustomHeaders: item.CustomHeaders,
		Status:        item.Status,
	}
}

func normalizeTunnel(t *store.Tunnel) error {
	if t.ListenIP == "" {
		t.ListenIP = "0.0.0.0"
	}
	if t.Mode == "" {
		t.Mode = "tcp"
	}
	if t.ListenPort <= 0 || t.TargetAddr == "" {
		return errImportField("listen_port and target_addr required")
	}
	if t.IPForwardMode == "" {
		t.IPForwardMode = "replace"
	}
	return nil
}

func normalizeHost(h *store.Host) error {
	if h.Host == "" || h.TargetAddr == "" {
		return errImportField("host and target_addr required")
	}
	if h.Location == "" {
		h.Location = "/"
	}
	if h.Scheme == "" {
		h.Scheme = "all"
	}
	if h.IPForwardMode == "" {
		h.IPForwardMode = "replace"
	}
	return nil
}

type importFieldError string

func errImportField(msg string) error { return importFieldError(msg) }

func (e importFieldError) Error() string { return string(e) }

func formatImportError(kind string, index int, err error) string {
	return kind + " #" + strconv.Itoa(index) + ": " + err.Error()
}

func tunnelClientNames(st *store.Store, tunnels []store.Tunnel) map[uint]string {
	ids := make(map[uint]struct{})
	for _, t := range tunnels {
		ids[t.ClientID] = struct{}{}
	}
	names := make(map[uint]string, len(ids))
	for id := range ids {
		if c, err := st.GetClientByID(id); err == nil {
			names[id] = c.Name
		}
	}
	return names
}

func hostClientNames(st *store.Store, hosts []store.Host) map[uint]string {
	ids := make(map[uint]struct{})
	for _, h := range hosts {
		ids[h.ClientID] = struct{}{}
	}
	names := make(map[uint]string, len(ids))
	for id := range ids {
		if c, err := st.GetClientByID(id); err == nil {
			names[id] = c.Name
		}
	}
	return names
}