server.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883
  1. package raft
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "fmt"
  6. "io"
  7. "net"
  8. "net/http"
  9. "sort"
  10. "strings"
  11. "sync"
  12. "time"
  13. "igit.com/xbase/raft/db"
  14. )
  15. // KVServer wraps Raft to provide a distributed key-value store
  16. type KVServer struct {
  17. Raft *Raft
  18. DB *db.Engine
  19. CLI *CLI
  20. AuthManager *AuthManager
  21. Watcher *WebHookWatcher
  22. httpServer *http.Server
  23. stopCh chan struct{}
  24. wg sync.WaitGroup
  25. stopOnce sync.Once
  26. // leavingNodes tracks nodes that are currently being removed
  27. // to prevent auto-rejoin/discovery logic from interfering
  28. leavingNodes sync.Map
  29. }
  30. // NewKVServer creates a new KV server
  31. func NewKVServer(config *Config) (*KVServer, error) {
  32. // Initialize DB Engine
  33. // Use a subdirectory for DB to avoid conflict with Raft logs if they share DataDir
  34. dbPath := config.DataDir + "/kv_engine"
  35. engine, err := db.NewEngine(dbPath)
  36. if err != nil {
  37. return nil, fmt.Errorf("failed to create db engine: %w", err)
  38. }
  39. // Initialize LastAppliedIndex from DB to prevent re-applying entries
  40. config.LastAppliedIndex = engine.GetLastAppliedIndex()
  41. // Create stop channel early for use in callbacks
  42. stopCh := make(chan struct{})
  43. // Configure snapshot provider
  44. config.SnapshotProvider = func(minIncludeIndex uint64) ([]byte, error) {
  45. // Wait for DB to catch up to the requested index
  46. // This is critical for data integrity during compaction
  47. for engine.GetLastAppliedIndex() < minIncludeIndex {
  48. select {
  49. case <-stopCh:
  50. return nil, fmt.Errorf("server stopping")
  51. default:
  52. time.Sleep(10 * time.Millisecond)
  53. }
  54. }
  55. // Force sync to disk to ensure data durability before compaction
  56. // This prevents data loss if Raft logs are compacted but DB data is only in OS cache
  57. if err := engine.Sync(); err != nil {
  58. return nil, fmt.Errorf("failed to sync engine before snapshot: %w", err)
  59. }
  60. return engine.Snapshot()
  61. }
  62. // Configure get handler for remote reads
  63. config.GetHandler = func(key string) (string, bool) {
  64. return engine.Get(key)
  65. }
  66. applyCh := make(chan ApplyMsg, 1000) // Increase buffer for async processing
  67. transport := NewTCPTransport(config.ListenAddr, 10, config.Logger)
  68. // Initialize WebHookWatcher
  69. // 5 workers, 3 retries
  70. watcher := NewWebHookWatcher(5, 3, config.Logger)
  71. r, err := NewRaft(config, transport, applyCh)
  72. if err != nil {
  73. engine.Close()
  74. return nil, err
  75. }
  76. s := &KVServer{
  77. Raft: r,
  78. DB: engine,
  79. CLI: nil,
  80. AuthManager: nil, // initialized below
  81. Watcher: watcher,
  82. stopCh: stopCh,
  83. }
  84. // Initialize AuthManager
  85. s.AuthManager = NewAuthManager(s)
  86. // Initialize CLI
  87. s.CLI = NewCLI(s)
  88. // Start applying entries
  89. go s.runApplyLoop(applyCh)
  90. // Start background maintenance loop
  91. s.wg.Add(1)
  92. go s.maintenanceLoop()
  93. return s, nil
  94. }
  95. func (s *KVServer) Start() error {
  96. // Start CLI if enabled
  97. if s.Raft.config.EnableCLI {
  98. go s.CLI.Start()
  99. }
  100. // Start HTTP Server if configured
  101. if s.Raft.config.HTTPAddr != "" {
  102. if err := s.startHTTPServer(s.Raft.config.HTTPAddr); err != nil {
  103. s.Raft.config.Logger.Warn("Failed to start HTTP server: %v", err)
  104. }
  105. }
  106. return s.Raft.Start()
  107. }
  108. func (s *KVServer) Stop() error {
  109. var err error
  110. s.stopOnce.Do(func() {
  111. // Stop maintenance loop
  112. if s.stopCh != nil {
  113. close(s.stopCh)
  114. s.wg.Wait()
  115. }
  116. // Stop Watcher
  117. if s.Watcher != nil {
  118. s.Watcher.Stop()
  119. }
  120. // Stop HTTP Server
  121. if s.httpServer != nil {
  122. s.httpServer.Close()
  123. }
  124. // Stop Raft first
  125. if errRaft := s.Raft.Stop(); errRaft != nil {
  126. err = errRaft
  127. }
  128. // Close DB
  129. if s.DB != nil {
  130. if errDB := s.DB.Close(); errDB != nil {
  131. // Combine errors if both fail
  132. if err != nil {
  133. err = fmt.Errorf("raft stop error: %v, db close error: %v", err, errDB)
  134. } else {
  135. err = errDB
  136. }
  137. }
  138. }
  139. })
  140. return err
  141. }
  142. func (s *KVServer) runApplyLoop(applyCh chan ApplyMsg) {
  143. for msg := range applyCh {
  144. if msg.CommandValid {
  145. // Optimization: Skip if already applied
  146. // We check this here to avoid unmarshalling and locking DB for known duplicates
  147. if msg.CommandIndex <= s.DB.GetLastAppliedIndex() {
  148. continue
  149. }
  150. var cmd KVCommand
  151. if err := json.Unmarshal(msg.Command, &cmd); err != nil {
  152. s.Raft.config.Logger.Error("Failed to unmarshal command: %v", err)
  153. continue
  154. }
  155. var err error
  156. switch cmd.Type {
  157. case KVSet:
  158. // Update Auth Cache for system keys
  159. if strings.HasPrefix(cmd.Key, SystemKeyPrefix) {
  160. s.AuthManager.UpdateCache(cmd.Key, cmd.Value, false)
  161. }
  162. err = s.DB.Set(cmd.Key, cmd.Value, msg.CommandIndex)
  163. if err == nil {
  164. s.Watcher.Notify(cmd.Key, cmd.Value, KVSet)
  165. }
  166. case KVDel:
  167. // Update Auth Cache for system keys
  168. if strings.HasPrefix(cmd.Key, SystemKeyPrefix) {
  169. s.AuthManager.UpdateCache(cmd.Key, "", true)
  170. }
  171. err = s.DB.Delete(cmd.Key, msg.CommandIndex)
  172. if err == nil {
  173. s.Watcher.Notify(cmd.Key, "", KVDel)
  174. }
  175. default:
  176. s.Raft.config.Logger.Error("Unknown command type: %d", cmd.Type)
  177. }
  178. if err != nil {
  179. s.Raft.config.Logger.Error("DB Apply failed: %v", err)
  180. }
  181. } else if msg.SnapshotValid {
  182. if err := s.DB.Restore(msg.Snapshot); err != nil {
  183. s.Raft.config.Logger.Error("DB Restore failed: %v", err)
  184. }
  185. }
  186. }
  187. }
  188. // SetAuthenticated sets a key-value pair with permission check
  189. func (s *KVServer) SetAuthenticated(key, value, token string) error {
  190. if err := s.AuthManager.CheckPermission(token, key, ActionWrite, value); err != nil {
  191. return err
  192. }
  193. return s.Set(key, value)
  194. }
  195. // DelAuthenticated deletes a key with permission check
  196. func (s *KVServer) DelAuthenticated(key, token string) error {
  197. if err := s.AuthManager.CheckPermission(token, key, ActionWrite, ""); err != nil {
  198. return err
  199. }
  200. return s.Del(key)
  201. }
  202. // Set sets a key-value pair
  203. func (s *KVServer) Set(key, value string) error {
  204. cmd := KVCommand{
  205. Type: KVSet,
  206. Key: key,
  207. Value: value,
  208. }
  209. data, err := json.Marshal(cmd)
  210. if err != nil {
  211. return err
  212. }
  213. _, _, err = s.Raft.ProposeWithForward(data)
  214. return err
  215. }
  216. // Del deletes a key
  217. func (s *KVServer) Del(key string) error {
  218. cmd := KVCommand{
  219. Type: KVDel,
  220. Key: key,
  221. }
  222. data, err := json.Marshal(cmd)
  223. if err != nil {
  224. return err
  225. }
  226. _, _, err = s.Raft.ProposeWithForward(data)
  227. return err
  228. }
  229. // GetAuthenticated gets a value with permission check (local read)
  230. func (s *KVServer) GetAuthenticated(key, token string) (string, bool, error) {
  231. if err := s.AuthManager.CheckPermission(token, key, ActionRead, ""); err != nil {
  232. return "", false, err
  233. }
  234. val, ok := s.Get(key)
  235. return val, ok, nil
  236. }
  237. // Get gets a value (local read, can be stale)
  238. // For linearizable reads, use GetLinear instead
  239. func (s *KVServer) Get(key string) (string, bool) {
  240. return s.DB.Get(key)
  241. }
  242. // GetLinearAuthenticated gets a value with linearizable consistency and permission check
  243. func (s *KVServer) GetLinearAuthenticated(key, token string) (string, bool, error) {
  244. if err := s.AuthManager.CheckPermission(token, key, ActionRead, ""); err != nil {
  245. return "", false, err
  246. }
  247. return s.GetLinear(key)
  248. }
  249. // Logout invalidates the current session
  250. func (s *KVServer) Logout(token string) error {
  251. return s.AuthManager.Logout(token)
  252. }
  253. // GetSessionInfo returns the session details
  254. func (s *KVServer) GetSessionInfo(token string) (*Session, error) {
  255. return s.AuthManager.GetSession(token)
  256. }
  257. // IsRoot checks if the token belongs to the root user
  258. func (s *KVServer) IsRoot(token string) bool {
  259. sess, err := s.AuthManager.GetSession(token)
  260. if err != nil {
  261. return false
  262. }
  263. return sess.Username == "root"
  264. }
  265. // SearchAuthenticated searches keys with permission checks
  266. func (s *KVServer) SearchAuthenticated(pattern string, limit, offset int, token string) ([]db.QueryResult, error) {
  267. // Optimization: If user has full access (*), delegate to DB with limit/offset
  268. if s.AuthManager.HasFullAccess(token) {
  269. sql := fmt.Sprintf("key like \"%s\" LIMIT %d OFFSET %d", pattern, limit, offset)
  270. return s.DB.Query(sql)
  271. }
  272. // Slow path: Fetch all potential matches and filter
  273. // We construct a SQL query that retrieves CANDIDATES.
  274. // We cannot safely apply LIMIT/OFFSET in SQL because we might filter out some results.
  275. // So we must fetch ALL matches.
  276. // WARNING: This can be slow and memory intensive.
  277. // If pattern is wildcard, we might fetch everything.
  278. sql := fmt.Sprintf("key like \"%s\"", pattern)
  279. results, err := s.DB.Query(sql)
  280. if err != nil {
  281. return nil, err
  282. }
  283. filtered := make([]db.QueryResult, 0, len(results))
  284. // Apply Filtering
  285. for _, r := range results {
  286. if err := s.AuthManager.CheckPermission(token, r.Key, ActionRead, ""); err == nil {
  287. filtered = append(filtered, r)
  288. }
  289. }
  290. // Apply Pagination on Filtered Results
  291. if offset > len(filtered) {
  292. return []db.QueryResult{}, nil
  293. }
  294. start := offset
  295. end := offset + limit
  296. if end > len(filtered) {
  297. end = len(filtered)
  298. }
  299. return filtered[start:end], nil
  300. }
  301. // CountAuthenticated counts keys with permission checks
  302. func (s *KVServer) CountAuthenticated(pattern string, token string) (int, error) {
  303. // Optimization: If user has full access
  304. if s.AuthManager.HasFullAccess(token) {
  305. sql := ""
  306. if pattern == "*" {
  307. sql = "*"
  308. } else {
  309. sql = fmt.Sprintf("key like \"%s\"", pattern)
  310. }
  311. return s.DB.Count(sql)
  312. }
  313. // Slow path: Iterate and check
  314. // We use DB.Query to get keys (Query returns values too, which is wasteful but DB engine doesn't expose keys-only Query yet)
  315. // Actually, we can access s.DB.Index.WalkPrefix if we want to be faster and avoid value read,
  316. // but s.DB.Index is inside 'db' package. We can access it if it's exported.
  317. // 'db' package exports 'Engine' and 'FlatIndex'.
  318. // So s.DB.Index IS accessible.
  319. count := 0
  320. // Determine prefix from pattern
  321. prefix := ""
  322. if strings.HasSuffix(pattern, "*") {
  323. prefix = strings.TrimSuffix(pattern, "*")
  324. } else if pattern == "*" {
  325. prefix = ""
  326. } else {
  327. // Exact match check
  328. if err := s.AuthManager.CheckPermission(token, pattern, ActionRead, ""); err == nil {
  329. // Check if exists
  330. if _, ok := s.DB.Get(pattern); ok {
  331. return 1, nil
  332. }
  333. return 0, nil
  334. }
  335. return 0, nil // No perm or not found
  336. }
  337. // Walk
  338. // Note: WalkPrefix locks the DB Index (Read Lock).
  339. // Calling CheckPermission inside might involve some logic but it is memory-only and usually fast.
  340. // However, if CheckPermission takes time, we hold DB lock.
  341. s.DB.Index.WalkPrefix(prefix, func(key string, entry db.IndexEntry) bool {
  342. // Check pattern match first (WalkPrefix is just prefix, pattern might be more complex like "user.*.name")
  343. if !db.WildcardMatch(key, pattern) {
  344. return true
  345. }
  346. if err := s.AuthManager.CheckPermission(token, key, ActionRead, ""); err == nil {
  347. count++
  348. }
  349. return true
  350. })
  351. return count, nil
  352. }
  353. // GetLinear gets a value with linearizable consistency
  354. // This ensures the read sees all writes committed before the read started
  355. func (s *KVServer) GetLinear(key string) (string, bool, error) {
  356. // First, ensure we have up-to-date data via ReadIndex
  357. _, err := s.Raft.ReadIndex()
  358. if err != nil {
  359. // If we're not leader, try forwarding
  360. if errors.Is(err, ErrNotLeader) {
  361. return s.forwardGet(key)
  362. }
  363. return "", false, err
  364. }
  365. val, ok := s.DB.Get(key)
  366. return val, ok, nil
  367. }
  368. // forwardGet forwards a get request to the leader
  369. func (s *KVServer) forwardGet(key string) (string, bool, error) {
  370. return s.Raft.ForwardGet(key)
  371. }
  372. // Join joins an existing cluster
  373. func (s *KVServer) Join(nodeID, addr string) error {
  374. return s.Raft.AddNodeWithForward(nodeID, addr)
  375. }
  376. // Leave leaves the cluster
  377. func (s *KVServer) Leave(nodeID string) error {
  378. // Mark node as leaving to prevent auto-rejoin
  379. s.leavingNodes.Store(nodeID, time.Now())
  380. // Auto-expire the leaving flag after a while
  381. go func() {
  382. time.Sleep(30 * time.Second)
  383. s.leavingNodes.Delete(nodeID)
  384. }()
  385. // Remove from RaftNode discovery key first to prevent auto-rejoin
  386. if err := s.removeNodeFromDiscovery(nodeID); err != nil {
  387. s.Raft.config.Logger.Warn("Failed to remove node from discovery key: %v", err)
  388. // Continue anyway, as the main goal is to leave the cluster
  389. }
  390. return s.Raft.RemoveNodeWithForward(nodeID)
  391. }
  392. // removeNodeFromDiscovery removes a node from the RaftNode key to prevent auto-rejoin
  393. func (s *KVServer) removeNodeFromDiscovery(targetID string) error {
  394. val, ok := s.Get("RaftNode")
  395. if !ok || val == "" {
  396. return nil
  397. }
  398. parts := strings.Split(val, ";")
  399. var newParts []string
  400. changed := false
  401. for _, part := range parts {
  402. if part == "" {
  403. continue
  404. }
  405. kv := strings.SplitN(part, "=", 2)
  406. if len(kv) == 2 {
  407. if kv[0] == targetID {
  408. changed = true
  409. continue // Skip this node
  410. }
  411. newParts = append(newParts, part)
  412. }
  413. }
  414. if changed {
  415. newVal := strings.Join(newParts, ";")
  416. return s.Set("RaftNode", newVal)
  417. }
  418. return nil
  419. }
  420. // WaitForLeader waits until a leader is elected
  421. func (s *KVServer) WaitForLeader(timeout time.Duration) error {
  422. deadline := time.Now().Add(timeout)
  423. for time.Now().Before(deadline) {
  424. leader := s.Raft.GetLeaderID()
  425. if leader != "" {
  426. return nil
  427. }
  428. time.Sleep(100 * time.Millisecond)
  429. }
  430. return fmt.Errorf("timeout waiting for leader")
  431. }
  432. // HealthCheck returns the health status of this server
  433. func (s *KVServer) HealthCheck() HealthStatus {
  434. return s.Raft.HealthCheck()
  435. }
  436. // GetStats returns runtime statistics
  437. func (s *KVServer) GetStats() Stats {
  438. return s.Raft.GetStats()
  439. }
  440. // GetMetrics returns runtime metrics
  441. func (s *KVServer) GetMetrics() Metrics {
  442. return s.Raft.GetMetrics()
  443. }
  444. // TransferLeadership transfers leadership to the specified node
  445. func (s *KVServer) TransferLeadership(targetID string) error {
  446. return s.Raft.TransferLeadership(targetID)
  447. }
  448. // GetClusterNodes returns current cluster membership
  449. func (s *KVServer) GetClusterNodes() map[string]string {
  450. return s.Raft.GetClusterNodes()
  451. }
  452. // IsLeader returns true if this node is the leader
  453. func (s *KVServer) IsLeader() bool {
  454. _, isLeader := s.Raft.GetState()
  455. return isLeader
  456. }
  457. // GetLeaderID returns the current leader ID
  458. func (s *KVServer) GetLeaderID() string {
  459. return s.Raft.GetLeaderID()
  460. }
  461. // GetLogSize returns the raft log size
  462. func (s *KVServer) GetLogSize() int64 {
  463. return s.Raft.log.GetLogSize()
  464. }
  465. // GetDBSize returns the db size
  466. func (s *KVServer) GetDBSize() int64 {
  467. return s.DB.GetDBSize()
  468. }
  469. // WatchURL registers a webhook url for a key
  470. func (s *KVServer) WatchURL(key, url string) {
  471. s.Watcher.Subscribe(key, url)
  472. }
  473. // UnwatchURL removes a webhook url for a key
  474. func (s *KVServer) UnwatchURL(key, url string) {
  475. s.Watcher.Unsubscribe(key, url)
  476. }
  477. // WatchAll registers a watcher for all keys
  478. func (s *KVServer) WatchAll(handler WatchHandler) {
  479. // s.FSM.WatchAll(handler)
  480. // TODO: Implement Watcher for DB
  481. }
  482. // Watch registers a watcher for a key
  483. func (s *KVServer) Watch(key string, handler WatchHandler) {
  484. // s.FSM.Watch(key, handler)
  485. // TODO: Implement Watcher for DB
  486. }
  487. // Unwatch removes watchers for a key
  488. func (s *KVServer) Unwatch(key string) {
  489. // s.FSM.Unwatch(key)
  490. // TODO: Implement Watcher for DB
  491. }
  492. func (s *KVServer) maintenanceLoop() {
  493. defer s.wg.Done()
  494. // Check every 1 second for faster reaction
  495. ticker := time.NewTicker(1 * time.Second)
  496. defer ticker.Stop()
  497. for {
  498. select {
  499. case <-s.stopCh:
  500. return
  501. case <-ticker.C:
  502. s.updateNodeInfo()
  503. s.checkConnections()
  504. }
  505. }
  506. }
  507. func (s *KVServer) updateNodeInfo() {
  508. // 1. Ensure "CreateNode/<NodeID>" is set to self address
  509. // We do this via Propose (Set) so it's replicated
  510. myID := s.Raft.config.NodeID
  511. myAddr := s.Raft.config.ListenAddr
  512. key := fmt.Sprintf("CreateNode/%s", myID)
  513. // Check if we need to update (avoid spamming logs/proposals)
  514. val, exists := s.Get(key)
  515. if !exists || val != myAddr {
  516. // Run in goroutine to avoid blocking
  517. go func() {
  518. if err := s.Set(key, myAddr); err != nil {
  519. s.Raft.config.Logger.Debug("Failed to update node info: %v", err)
  520. }
  521. }()
  522. }
  523. // 2. Only leader updates RaftNode aggregation
  524. if s.IsLeader() {
  525. // Read current RaftNode to preserve history
  526. currentVal, _ := s.Get("RaftNode")
  527. knownNodes := make(map[string]string)
  528. if currentVal != "" {
  529. parts := strings.Split(currentVal, ";")
  530. for _, part := range parts {
  531. if part == "" { continue }
  532. kv := strings.SplitN(part, "=", 2)
  533. if len(kv) == 2 {
  534. knownNodes[kv[0]] = kv[1]
  535. }
  536. }
  537. }
  538. // Merge current cluster nodes
  539. changed := false
  540. currentCluster := s.GetClusterNodes()
  541. for id, addr := range currentCluster {
  542. // Skip nodes that are marked as leaving
  543. if _, leaving := s.leavingNodes.Load(id); leaving {
  544. continue
  545. }
  546. if knownNodes[id] != addr {
  547. knownNodes[id] = addr
  548. changed = true
  549. }
  550. }
  551. // If changed, update RaftNode
  552. if changed {
  553. var peers []string
  554. for id, addr := range knownNodes {
  555. peers = append(peers, fmt.Sprintf("%s=%s", id, addr))
  556. }
  557. sort.Strings(peers)
  558. newVal := strings.Join(peers, ";")
  559. // Check again if we need to write to avoid loops if Get returned stale
  560. if newVal != currentVal {
  561. go func(k, v string) {
  562. if err := s.Set(k, v); err != nil {
  563. s.Raft.config.Logger.Warn("Failed to update RaftNode key: %v", err)
  564. }
  565. }("RaftNode", newVal)
  566. }
  567. }
  568. }
  569. }
  570. func (s *KVServer) checkConnections() {
  571. if !s.IsLeader() {
  572. return
  573. }
  574. // Read RaftNode key to find potential members that are missing
  575. val, ok := s.Get("RaftNode")
  576. if !ok || val == "" {
  577. return
  578. }
  579. // Parse saved nodes
  580. savedParts := strings.Split(val, ";")
  581. currentNodes := s.GetClusterNodes()
  582. // Invert currentNodes for address check
  583. currentAddrs := make(map[string]bool)
  584. for _, addr := range currentNodes {
  585. currentAddrs[addr] = true
  586. }
  587. for _, part := range savedParts {
  588. if part == "" {
  589. continue
  590. }
  591. // Expect id=addr
  592. kv := strings.SplitN(part, "=", 2)
  593. if len(kv) != 2 {
  594. continue
  595. }
  596. id, addr := kv[0], kv[1]
  597. // Skip invalid addresses
  598. if strings.HasPrefix(addr, ".") || !strings.Contains(addr, ":") {
  599. continue
  600. }
  601. if !currentAddrs[addr] {
  602. // Skip nodes that are marked as leaving
  603. if _, leaving := s.leavingNodes.Load(id); leaving {
  604. continue
  605. }
  606. // Found a node that was previously in the cluster but is now missing
  607. // Try to add it back
  608. // We use AddNodeWithForward which handles non-blocking internally somewhat,
  609. // but we should run this in goroutine to not block the loop
  610. go func(nodeID, nodeAddr string) {
  611. // Try to add node
  612. s.Raft.config.Logger.Info("Auto-rejoining node found in RaftNode: %s (%s)", nodeID, nodeAddr)
  613. if err := s.Join(nodeID, nodeAddr); err != nil {
  614. s.Raft.config.Logger.Debug("Failed to auto-rejoin node %s: %v", nodeID, err)
  615. }
  616. }(id, addr)
  617. }
  618. }
  619. }
  620. // startHTTPServer starts the HTTP API server
  621. func (s *KVServer) startHTTPServer(addr string) error {
  622. mux := http.NewServeMux()
  623. // KV API
  624. mux.HandleFunc("/kv", func(w http.ResponseWriter, r *http.Request) {
  625. token := r.Header.Get("X-Raft-Token")
  626. switch r.Method {
  627. case http.MethodGet:
  628. key := r.URL.Query().Get("key")
  629. if key == "" {
  630. http.Error(w, "missing key", http.StatusBadRequest)
  631. return
  632. }
  633. // Use Authenticated method
  634. val, found, err := s.GetLinearAuthenticated(key, token)
  635. if err != nil {
  636. // Distinguish auth error vs raft error?
  637. if strings.Contains(err.Error(), "permission") || strings.Contains(err.Error(), "unauthorized") {
  638. http.Error(w, err.Error(), http.StatusForbidden)
  639. } else {
  640. http.Error(w, err.Error(), http.StatusInternalServerError)
  641. }
  642. return
  643. }
  644. if !found {
  645. http.Error(w, "not found", http.StatusNotFound)
  646. return
  647. }
  648. w.Write([]byte(val))
  649. case http.MethodPost:
  650. body, _ := io.ReadAll(r.Body)
  651. var req struct {
  652. Key string `json:"key"`
  653. Value string `json:"value"`
  654. }
  655. if err := json.Unmarshal(body, &req); err != nil {
  656. http.Error(w, "invalid json", http.StatusBadRequest)
  657. return
  658. }
  659. if err := s.SetAuthenticated(req.Key, req.Value, token); err != nil {
  660. if strings.Contains(err.Error(), "permission") || strings.Contains(err.Error(), "unauthorized") {
  661. http.Error(w, err.Error(), http.StatusForbidden)
  662. } else {
  663. http.Error(w, err.Error(), http.StatusInternalServerError)
  664. }
  665. return
  666. }
  667. w.WriteHeader(http.StatusOK)
  668. case http.MethodDelete:
  669. key := r.URL.Query().Get("key")
  670. if key == "" {
  671. http.Error(w, "missing key", http.StatusBadRequest)
  672. return
  673. }
  674. if err := s.DelAuthenticated(key, token); err != nil {
  675. if strings.Contains(err.Error(), "permission") || strings.Contains(err.Error(), "unauthorized") {
  676. http.Error(w, err.Error(), http.StatusForbidden)
  677. } else {
  678. http.Error(w, err.Error(), http.StatusInternalServerError)
  679. }
  680. return
  681. }
  682. w.WriteHeader(http.StatusOK)
  683. default:
  684. http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
  685. }
  686. })
  687. // Auth API
  688. mux.HandleFunc("/auth/login", func(w http.ResponseWriter, r *http.Request) {
  689. if r.Method != http.MethodPost {
  690. http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
  691. return
  692. }
  693. var req struct {
  694. Username string `json:"username"`
  695. Password string `json:"password"`
  696. Code string `json:"code"`
  697. }
  698. body, _ := io.ReadAll(r.Body)
  699. if err := json.Unmarshal(body, &req); err != nil {
  700. http.Error(w, "invalid json", http.StatusBadRequest)
  701. return
  702. }
  703. ip := r.RemoteAddr
  704. if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
  705. ip = host
  706. }
  707. token, err := s.AuthManager.Login(req.Username, req.Password, req.Code, ip)
  708. if err != nil {
  709. http.Error(w, err.Error(), http.StatusUnauthorized)
  710. return
  711. }
  712. resp := struct {
  713. Token string `json:"token"`
  714. }{Token: token}
  715. json.NewEncoder(w).Encode(resp)
  716. })
  717. // Watcher API
  718. mux.HandleFunc("/watch", func(w http.ResponseWriter, r *http.Request) {
  719. if r.Method != http.MethodPost {
  720. http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
  721. return
  722. }
  723. body, _ := io.ReadAll(r.Body)
  724. var req struct {
  725. Key string `json:"key"`
  726. URL string `json:"url"`
  727. }
  728. if err := json.Unmarshal(body, &req); err != nil {
  729. http.Error(w, "invalid json", http.StatusBadRequest)
  730. return
  731. }
  732. if req.Key == "" || req.URL == "" {
  733. http.Error(w, "missing key or url", http.StatusBadRequest)
  734. return
  735. }
  736. s.WatchURL(req.Key, req.URL)
  737. w.WriteHeader(http.StatusOK)
  738. })
  739. mux.HandleFunc("/unwatch", func(w http.ResponseWriter, r *http.Request) {
  740. if r.Method != http.MethodPost {
  741. http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
  742. return
  743. }
  744. body, _ := io.ReadAll(r.Body)
  745. var req struct {
  746. Key string `json:"key"`
  747. URL string `json:"url"`
  748. }
  749. if err := json.Unmarshal(body, &req); err != nil {
  750. http.Error(w, "invalid json", http.StatusBadRequest)
  751. return
  752. }
  753. s.UnwatchURL(req.Key, req.URL)
  754. w.WriteHeader(http.StatusOK)
  755. })
  756. s.httpServer = &http.Server{
  757. Addr: addr,
  758. Handler: mux,
  759. }
  760. go func() {
  761. s.Raft.config.Logger.Info("HTTP API server listening on %s", addr)
  762. if err := s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
  763. s.Raft.config.Logger.Error("HTTP server failed: %v", err)
  764. }
  765. }()
  766. return nil
  767. }