goc/pkg/server/server.go

/*
 Copyright 2021 Qiniu Cloud (qiniu.com)
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
     http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package server

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"net/rpc"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/qiniu/goc/v2/pkg/log"
	"github.com/qiniu/goc/v2/pkg/server/store"
)

// gocServer represents a goc server
type gocServer struct {
	port  int
	store store.Store

	upgrader websocket.Upgrader

	agents sync.Map

	watchCh      chan []byte
	watchClients sync.Map

	idCount int64
}

type gocCliendId string

const (
	DISCONNECT   = 1 << iota
	RPCCONNECT   = 1 << iota
	WATCHCONNECT = 1 << iota
)

// gocCoveredAgent represents a covered client
type gocCoveredAgent struct {
	Id            string `json:"id"`
	RpcRemoteIP   string `json:"rpc_remoteip"`
	WatchRemoteIP string `json:"watch_remoteip"`
	Hostname      string `json:"hostname"`
	CmdLine       string `json:"cmdline"`
	Pid           string `json:"pid"`

	// 用户可以选择上报一些定制信息
	// 比如不同 namespace 的 statefulset POD，它们的 hostname/cmdline/pid 都是一样的，
	// 这时候将 extra 设置为 namespace 并上报，这个额外的信息在展示时将更友好
	Extra string `json:"extra"`

	Token  string `json:"token"`
	Status int    `json:"status"` // 表示该 agent 是否处于 connected 状态

	rpc *rpc.Client `json:"-"`

	exitCh             chan int `json:"-"`
	closeRpcConnOnce   func()   `json:"-"` // close rpc conn 只执行一次
	closeWatchConnOnce func()   `json:"-"` // close watch conn 只执行一次
}

func (agent *gocCoveredAgent) closeConnection() {
	if agent.closeRpcConnOnce != nil {
		agent.closeRpcConnOnce()
	}

	if agent.closeWatchConnOnce != nil {
		agent.closeWatchConnOnce()
	}
}

//  api 客户端，不是 agent
type gocWatchClient struct {
	ws     *websocket.Conn
	exitCh chan int
	once   sync.Once
}

func RunGocServerUntilExit(host string, path string) {
	s, err := store.NewFileStore(path)
	if err != nil {
		log.Fatalf("cannot create store for goc server: %v", err)
	}

	gs := gocServer{
		store: s,
		upgrader: websocket.Upgrader{
			ReadBufferSize:   4096,
			WriteBufferSize:  4096,
			HandshakeTimeout: 45 * time.Second,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		watchCh: make(chan []byte, 4096),
	}

	// 从持久化存储上恢复 agent 列表
	gs.restoreAgents()

	r := gin.Default()
	v2 := r.Group("/v2")
	{
		v2.GET("/cover/profile", gs.getProfiles)
		v2.DELETE("/cover/profile", gs.resetProfiles)
		v2.GET("/agents", gs.listAgents)
		v2.DELETE("/agents", gs.removeAgents)

		v2.GET("/cover/ws/watch", gs.watchProfileUpdate)

		// internal use only
		v2.GET("/internal/register", gs.register)
		v2.GET("/internal/ws/rpcstream", gs.serveRpcStream)
		v2.GET("/internal/ws/watchstream", gs.serveWatchInternalStream)
	}

	go gs.watchLoop()

	r.Run(host)
}

func (gs *gocServer) register(c *gin.Context) {
	// 检查插桩服务上报的信息
	hostname := c.Query("hostname")
	pid := c.Query("pid")
	cmdline := c.Query("cmdline")
	extra := c.Query("extra")

	if hostname == "" || pid == "" || cmdline == "" {
		c.JSON(http.StatusBadRequest, gin.H{
			"msg": "missing some params",
		})
		return
	}

	atomic.AddInt64(&gs.idCount, 1)

	genToken := func(i int64) string {
		now := time.Now().UnixNano()
		random := rand.Int()

		raw := fmt.Sprintf("%v-%v-%v", i, random, now)
		sum := sha256.Sum256([]byte(raw))
		h := fmt.Sprintf("%x", sum[:16])

		return h
	}

	token := genToken(gs.idCount)
	id := strconv.Itoa(int(gs.idCount))

	agent := &gocCoveredAgent{
		Id:       id,
		Hostname: hostname,
		Pid:      pid,
		CmdLine:  cmdline,
		Token:    token,
		Status:   DISCONNECT,
		Extra:    extra,
	}

	// 持久化
	err := gs.saveAgentToStore(agent)
	if err != nil {
		log.Errorf("fail to save to store: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": err.Error(),
		})
	}
	// 维护 agent 连接
	gs.agents.Store(id, agent)

	log.Infof("one agent registered, id: %v, cmdline: %v, pid: %v, hostname: %v", id, agent.CmdLine, agent.Pid, agent.Hostname)

	c.JSON(http.StatusOK, gin.H{
		"id":    strconv.Itoa(int(gs.idCount)),
		"token": token,
	})
}

func (gs *gocServer) saveAgentToStore(agent *gocCoveredAgent) error {

	value, err := json.Marshal(agent)
	if err != nil {
		return err
	}
	return gs.store.Set("/goc/agents/"+agent.Id, string(value))
}

func (gs *gocServer) removeAgentFromStore(id string) error {

	return gs.store.Remove("/goc/agents/" + id)
}

func (gs *gocServer) removeAllAgentsFromStore() error {

	return gs.store.RangeRemove("/goc/agents/")
}

func (gs *gocServer) restoreAgents() {
	pattern := "/goc/agents/"

	// ignore err, 这个 err 不需要处理，直接忽略
	rawagents, _ := gs.store.Range(pattern)

	var maxId int
	for _, rawagent := range rawagents {
		var agent gocCoveredAgent
		err := json.Unmarshal([]byte(rawagent), &agent)
		if err != nil {
			log.Fatalf("fail to unmarshal restore agents: %v", err)
		}

		id, err := strconv.Atoi(agent.Id)
		if err != nil {
			log.Fatalf("fail to transfer id to number: %v", err)
		}
		if maxId < id {
			maxId = id
		}

		gs.agents.Store(agent.Id, &agent)
		log.Infof("restore one agent: %v, %v from store", id, agent.RpcRemoteIP)

		agent.RpcRemoteIP = ""
		agent.WatchRemoteIP = ""
		agent.Status = DISCONNECT
	}

	// 更新全局 id
	atomic.StoreInt64(&gs.idCount, int64(maxId))
}