mirror of
https://github.com/fankes/beszel.git
synced 2025-10-18 17:29:28 +08:00
322 lines
10 KiB
Go
322 lines
10 KiB
Go
package hub
|
|
|
|
import (
|
|
"errors"
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/henrygd/beszel/internal/common"
|
|
"github.com/henrygd/beszel/internal/hub/expirymap"
|
|
"github.com/henrygd/beszel/internal/hub/ws"
|
|
|
|
"github.com/blang/semver"
|
|
"github.com/lxzan/gws"
|
|
"github.com/pocketbase/dbx"
|
|
"github.com/pocketbase/pocketbase/core"
|
|
)
|
|
|
|
// agentConnectRequest holds information related to an agent's connection attempt.
|
|
type agentConnectRequest struct {
|
|
hub *Hub
|
|
req *http.Request
|
|
res http.ResponseWriter
|
|
token string
|
|
agentSemVer semver.Version
|
|
// isUniversalToken is true if the token is a universal token.
|
|
isUniversalToken bool
|
|
// userId is the user ID associated with the universal token.
|
|
userId string
|
|
}
|
|
|
|
// universalTokenMap stores active universal tokens and their associated user IDs.
|
|
var universalTokenMap tokenMap
|
|
|
|
type tokenMap struct {
|
|
store *expirymap.ExpiryMap[string]
|
|
once sync.Once
|
|
}
|
|
|
|
// getMap returns the expirymap, creating it if necessary.
|
|
func (tm *tokenMap) GetMap() *expirymap.ExpiryMap[string] {
|
|
tm.once.Do(func() {
|
|
tm.store = expirymap.New[string](time.Hour)
|
|
})
|
|
return tm.store
|
|
}
|
|
|
|
// handleAgentConnect is the HTTP handler for an agent's connection request.
|
|
func (h *Hub) handleAgentConnect(e *core.RequestEvent) error {
|
|
agentRequest := agentConnectRequest{req: e.Request, res: e.Response, hub: h}
|
|
_ = agentRequest.agentConnect()
|
|
return nil
|
|
}
|
|
|
|
// agentConnect validates agent credentials and upgrades the connection to a WebSocket.
|
|
func (acr *agentConnectRequest) agentConnect() (err error) {
|
|
var agentVersion string
|
|
|
|
acr.token, agentVersion, err = acr.validateAgentHeaders(acr.req.Header)
|
|
if err != nil {
|
|
return acr.sendResponseError(acr.res, http.StatusBadRequest, "")
|
|
}
|
|
|
|
// Check if token is an active universal token
|
|
acr.userId, acr.isUniversalToken = universalTokenMap.GetMap().GetOk(acr.token)
|
|
|
|
// Find matching fingerprint records for this token
|
|
fpRecords := getFingerprintRecordsByToken(acr.token, acr.hub)
|
|
if len(fpRecords) == 0 && !acr.isUniversalToken {
|
|
// Invalid token - no records found and not a universal token
|
|
return acr.sendResponseError(acr.res, http.StatusUnauthorized, "Invalid token")
|
|
}
|
|
|
|
// Validate agent version
|
|
acr.agentSemVer, err = semver.Parse(agentVersion)
|
|
if err != nil {
|
|
return acr.sendResponseError(acr.res, http.StatusUnauthorized, "Invalid agent version")
|
|
}
|
|
|
|
// Upgrade connection to WebSocket
|
|
conn, err := ws.GetUpgrader().Upgrade(acr.res, acr.req)
|
|
if err != nil {
|
|
return acr.sendResponseError(acr.res, http.StatusInternalServerError, "WebSocket upgrade failed")
|
|
}
|
|
|
|
go acr.verifyWsConn(conn, fpRecords)
|
|
|
|
return nil
|
|
}
|
|
|
|
// verifyWsConn verifies the WebSocket connection using the agent's fingerprint and
|
|
// SSH key signature, then adds the system to the system manager.
|
|
func (acr *agentConnectRequest) verifyWsConn(conn *gws.Conn, fpRecords []ws.FingerprintRecord) (err error) {
|
|
wsConn := ws.NewWsConnection(conn)
|
|
|
|
// must set wsConn in connection store before the read loop
|
|
conn.Session().Store("wsConn", wsConn)
|
|
|
|
// make sure connection is closed if there is an error
|
|
defer func() {
|
|
if err != nil {
|
|
wsConn.Close([]byte(err.Error()))
|
|
}
|
|
}()
|
|
|
|
go conn.ReadLoop()
|
|
|
|
signer, err := acr.hub.GetSSHKey("")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
agentFingerprint, err := wsConn.GetFingerprint(acr.token, signer, acr.isUniversalToken)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Find or create the appropriate system for this token and fingerprint
|
|
fpRecord, err := acr.findOrCreateSystemForToken(fpRecords, agentFingerprint)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return acr.hub.sm.AddWebSocketSystem(fpRecord.SystemId, acr.agentSemVer, wsConn)
|
|
}
|
|
|
|
// validateAgentHeaders extracts and validates the token and agent version from HTTP headers.
|
|
func (acr *agentConnectRequest) validateAgentHeaders(headers http.Header) (string, string, error) {
|
|
token := headers.Get("X-Token")
|
|
agentVersion := headers.Get("X-Beszel")
|
|
|
|
if agentVersion == "" || token == "" || len(token) > 64 {
|
|
return "", "", errors.New("")
|
|
}
|
|
return token, agentVersion, nil
|
|
}
|
|
|
|
// sendResponseError writes an HTTP error response.
|
|
func (acr *agentConnectRequest) sendResponseError(res http.ResponseWriter, code int, message string) error {
|
|
res.WriteHeader(code)
|
|
if message != "" {
|
|
res.Write([]byte(message))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// getFingerprintRecordsByToken retrieves all fingerprint records associated with a given token.
|
|
func getFingerprintRecordsByToken(token string, h *Hub) []ws.FingerprintRecord {
|
|
var records []ws.FingerprintRecord
|
|
// All will populate empty slice even on error
|
|
_ = h.DB().NewQuery("SELECT id, system, fingerprint, token FROM fingerprints WHERE token = {:token}").
|
|
Bind(dbx.Params{
|
|
"token": token,
|
|
}).
|
|
All(&records)
|
|
return records
|
|
}
|
|
|
|
// findOrCreateSystemForToken finds an existing system matching the token and fingerprint,
|
|
// or creates a new one for a universal token.
|
|
func (acr *agentConnectRequest) findOrCreateSystemForToken(fpRecords []ws.FingerprintRecord, agentFingerprint common.FingerprintResponse) (ws.FingerprintRecord, error) {
|
|
// No records - only valid for active universal tokens
|
|
if len(fpRecords) == 0 {
|
|
return acr.handleNoRecords(agentFingerprint)
|
|
}
|
|
|
|
// Single record - handle as regular token
|
|
if len(fpRecords) == 1 && !acr.isUniversalToken {
|
|
return acr.handleSingleRecord(fpRecords[0], agentFingerprint)
|
|
}
|
|
|
|
// Multiple records or universal token - look for matching fingerprint
|
|
return acr.handleMultipleRecordsOrUniversalToken(fpRecords, agentFingerprint)
|
|
}
|
|
|
|
// handleNoRecords handles the case where no fingerprint records are found for a token.
|
|
// A new system is created if the token is a valid universal token.
|
|
func (acr *agentConnectRequest) handleNoRecords(agentFingerprint common.FingerprintResponse) (ws.FingerprintRecord, error) {
|
|
var fpRecord ws.FingerprintRecord
|
|
|
|
if !acr.isUniversalToken || acr.userId == "" {
|
|
return fpRecord, errors.New("no matching fingerprints")
|
|
}
|
|
|
|
return acr.createNewSystemForUniversalToken(agentFingerprint)
|
|
}
|
|
|
|
// handleSingleRecord handles the case with a single fingerprint record. It validates
|
|
// the agent's fingerprint against the stored one, or sets it on first connect.
|
|
func (acr *agentConnectRequest) handleSingleRecord(fpRecord ws.FingerprintRecord, agentFingerprint common.FingerprintResponse) (ws.FingerprintRecord, error) {
|
|
// If no current fingerprint, update with new fingerprint (first time connecting)
|
|
if fpRecord.Fingerprint == "" {
|
|
if err := acr.hub.SetFingerprint(&fpRecord, agentFingerprint.Fingerprint); err != nil {
|
|
return fpRecord, err
|
|
}
|
|
// Update the record with the fingerprint that was set
|
|
fpRecord.Fingerprint = agentFingerprint.Fingerprint
|
|
return fpRecord, nil
|
|
}
|
|
|
|
// Abort if fingerprint exists but doesn't match (different machine)
|
|
if fpRecord.Fingerprint != agentFingerprint.Fingerprint {
|
|
return fpRecord, errors.New("fingerprint mismatch")
|
|
}
|
|
|
|
return fpRecord, nil
|
|
}
|
|
|
|
// handleMultipleRecordsOrUniversalToken finds a matching fingerprint from multiple records.
|
|
// If no match is found and the token is a universal token, a new system is created.
|
|
func (acr *agentConnectRequest) handleMultipleRecordsOrUniversalToken(fpRecords []ws.FingerprintRecord, agentFingerprint common.FingerprintResponse) (ws.FingerprintRecord, error) {
|
|
// Return existing record with matching fingerprint if found
|
|
for i := range fpRecords {
|
|
if fpRecords[i].Fingerprint == agentFingerprint.Fingerprint {
|
|
return fpRecords[i], nil
|
|
}
|
|
}
|
|
|
|
// No matching fingerprint record found, but it's
|
|
// an active universal token so create a new system
|
|
if acr.isUniversalToken {
|
|
return acr.createNewSystemForUniversalToken(agentFingerprint)
|
|
}
|
|
|
|
return ws.FingerprintRecord{}, errors.New("fingerprint mismatch")
|
|
}
|
|
|
|
// createNewSystemForUniversalToken creates a new system and fingerprint record for a universal token.
|
|
func (acr *agentConnectRequest) createNewSystemForUniversalToken(agentFingerprint common.FingerprintResponse) (ws.FingerprintRecord, error) {
|
|
var fpRecord ws.FingerprintRecord
|
|
if !acr.isUniversalToken || acr.userId == "" {
|
|
return fpRecord, errors.New("invalid token")
|
|
}
|
|
|
|
fpRecord.Token = acr.token
|
|
|
|
systemId, err := acr.createSystem(agentFingerprint)
|
|
if err != nil {
|
|
return fpRecord, err
|
|
}
|
|
fpRecord.SystemId = systemId
|
|
|
|
// Set the fingerprint for the new system
|
|
if err := acr.hub.SetFingerprint(&fpRecord, agentFingerprint.Fingerprint); err != nil {
|
|
return fpRecord, err
|
|
}
|
|
|
|
// Update the record with the fingerprint that was set
|
|
fpRecord.Fingerprint = agentFingerprint.Fingerprint
|
|
|
|
return fpRecord, nil
|
|
}
|
|
|
|
// createSystem creates a new system record in the database using details from the agent.
|
|
func (acr *agentConnectRequest) createSystem(agentFingerprint common.FingerprintResponse) (recordId string, err error) {
|
|
systemsCollection, err := acr.hub.FindCachedCollectionByNameOrId("systems")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
remoteAddr := getRealIP(acr.req)
|
|
// separate port from address
|
|
if agentFingerprint.Hostname == "" {
|
|
agentFingerprint.Hostname = remoteAddr
|
|
}
|
|
if agentFingerprint.Port == "" {
|
|
agentFingerprint.Port = "45876"
|
|
}
|
|
// create new record
|
|
systemRecord := core.NewRecord(systemsCollection)
|
|
systemRecord.Set("name", agentFingerprint.Hostname)
|
|
systemRecord.Set("host", remoteAddr)
|
|
systemRecord.Set("port", agentFingerprint.Port)
|
|
systemRecord.Set("users", []string{acr.userId})
|
|
|
|
return systemRecord.Id, acr.hub.Save(systemRecord)
|
|
}
|
|
|
|
// SetFingerprint creates or updates a fingerprint record in the database.
|
|
func (h *Hub) SetFingerprint(fpRecord *ws.FingerprintRecord, fingerprint string) (err error) {
|
|
// // can't use raw query here because it doesn't trigger SSE
|
|
var record *core.Record
|
|
switch fpRecord.Id {
|
|
case "":
|
|
// create new record for universal token
|
|
collection, _ := h.FindCachedCollectionByNameOrId("fingerprints")
|
|
record = core.NewRecord(collection)
|
|
record.Set("system", fpRecord.SystemId)
|
|
default:
|
|
record, err = h.FindRecordById("fingerprints", fpRecord.Id)
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
record.Set("token", fpRecord.Token)
|
|
record.Set("fingerprint", fingerprint)
|
|
return h.SaveNoValidate(record)
|
|
}
|
|
|
|
// getRealIP extracts the client's real IP address from request headers,
|
|
// checking common proxy headers before falling back to the remote address.
|
|
func getRealIP(r *http.Request) string {
|
|
if ip := r.Header.Get("CF-Connecting-IP"); ip != "" {
|
|
return ip
|
|
}
|
|
if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
|
|
// X-Forwarded-For can contain a comma-separated list: "client_ip, proxy1, proxy2"
|
|
// Take the first one
|
|
ips := strings.Split(ip, ",")
|
|
if len(ips) > 0 {
|
|
return strings.TrimSpace(ips[0])
|
|
}
|
|
}
|
|
// Fallback to RemoteAddr
|
|
ip, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err != nil {
|
|
return r.RemoteAddr
|
|
}
|
|
return ip
|
|
}
|