1  import Foundation
  2  
  3  /// Global socket path for signal handler cleanup (can't capture locals in @convention(c)).
  4  private var globalSocketPath: UnsafeMutablePointer<CChar>?
  5  
  6  private func signalHandler(_: Int32) {
  7      if let path = globalSocketPath {
  8          unlink(path)
  9          free(path)
 10      }
 11      exit(0)
 12  }
 13  
 14  /// Runs ax_server as a persistent Unix socket server.
 15  /// Same NDJSON protocol as the stdin/stdout mode — one JSON request per line,
 16  /// one JSON response per line. Accepts one client at a time; when the client
 17  /// disconnects, accepts the next connection.
 18  func runSocketServer(path socketPath: String) {
 19      // Clean up stale socket file
 20      unlink(socketPath)
 21  
 22      // Register signal handlers for clean shutdown
 23      globalSocketPath = strdup(socketPath)
 24      signal(SIGINT, signalHandler)
 25      signal(SIGTERM, signalHandler)
 26  
 27      // Create Unix domain socket
 28      let fd = socket(AF_UNIX, SOCK_STREAM, 0)
 29      guard fd >= 0 else {
 30          FileHandle.standardError.write("ax_server: failed to create socket\n".data(using: .utf8)!)
 31          exit(1)
 32      }
 33  
 34      var addr = sockaddr_un()
 35      addr.sun_family = sa_family_t(AF_UNIX)
 36      let pathBytes = socketPath.utf8CString
 37      guard pathBytes.count <= MemoryLayout.size(ofValue: addr.sun_path) else {
 38          FileHandle.standardError.write("ax_server: socket path too long\n".data(using: .utf8)!)
 39          exit(1)
 40      }
 41      withUnsafeMutablePointer(to: &addr.sun_path) { sunPath in
 42          sunPath.withMemoryRebound(to: CChar.self, capacity: pathBytes.count) { dst in
 43              for i in 0..<pathBytes.count {
 44                  dst[i] = pathBytes[i]
 45              }
 46          }
 47      }
 48  
 49      let bindResult = withUnsafePointer(to: &addr) { ptr in
 50          ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) { sockPtr in
 51              bind(fd, sockPtr, socklen_t(MemoryLayout<sockaddr_un>.size))
 52          }
 53      }
 54      guard bindResult == 0 else {
 55          FileHandle.standardError.write("ax_server: bind failed: \(String(cString: strerror(errno)))\n".data(using: .utf8)!)
 56          exit(1)
 57      }
 58  
 59      guard listen(fd, 1) == 0 else {
 60          FileHandle.standardError.write("ax_server: listen failed\n".data(using: .utf8)!)
 61          exit(1)
 62      }
 63  
 64      // Write "ready" to stdout so the launcher knows the socket is listening
 65      print("ready")
 66      fflush(stdout)
 67  
 68      let enc = JSONEncoder()
 69      enc.outputFormatting = [.sortedKeys]
 70      let dec = JSONDecoder()
 71  
 72      // Accept one client, serve it, then exit.
 73      // Each shan instance launches its own ax_server; there's no reason to
 74      // accept another connection after the client disconnects.
 75      let clientFD = accept(fd, nil, nil)
 76      guard clientFD >= 0 else {
 77          signalHandler(0)
 78          return
 79      }
 80  
 81      let input = FileHandle(fileDescriptor: clientFD, closeOnDealloc: false)
 82      let output = FileHandle(fileDescriptor: clientFD, closeOnDealloc: false)
 83  
 84      handleClient(input: input, output: output, encoder: enc, decoder: dec)
 85  
 86      close(clientFD)
 87      close(fd)
 88      // Client disconnected — clean up socket and exit
 89      signalHandler(0)
 90  }
 91  
 92  /// Process requests from a single client connection until it disconnects.
 93  private func handleClient(
 94      input: FileHandle,
 95      output: FileHandle,
 96      encoder: JSONEncoder,
 97      decoder: JSONDecoder
 98  ) {
 99      // Read data in chunks and split on newlines
100      var buffer = Data()
101  
102      while true {
103          let chunk = input.availableData
104          if chunk.isEmpty { break } // EOF — client disconnected
105  
106          buffer.append(chunk)
107  
108          // Process all complete lines in buffer
109          while let newlineIndex = buffer.firstIndex(of: UInt8(ascii: "\n")) {
110              let lineData = buffer[buffer.startIndex..<newlineIndex]
111              buffer = buffer[buffer.index(after: newlineIndex)...]
112  
113              guard !lineData.isEmpty else { continue }
114  
115              guard let req = try? decoder.decode(Request.self, from: Data(lineData)) else {
116                  let resp = Response(id: 0, error: ErrorInfo(code: -1, message: "Invalid JSON request"))
117                  writeToHandle(resp, encoder: encoder, output: output)
118                  continue
119              }
120  
121              let params = req.params ?? Params(
122                  pid: nil, maxDepth: nil, semanticBudget: nil, filter: nil,
123                  path: nil, expectedRole: nil, value: nil, appName: nil,
124                  query: nil, role: nil, identifier: nil, type: nil,
125                  x: nil, y: nil, button: nil, clicks: nil,
126                  key: nil, modifiers: nil, dx: nil, dy: nil,
127                  windowTitle: nil, verify: nil, condition: nil,
128                  timeout: nil, interval: nil, roles: nil, maxLabels: nil
129              )
130  
131              let response = dispatch(id: req.id, method: req.method, params: params)
132              writeToHandle(response, encoder: encoder, output: output)
133          }
134      }
135  }
136  
137  /// Write a JSON response followed by a newline to the given file handle.
138  private func writeToHandle(_ resp: Response, encoder: JSONEncoder, output: FileHandle) {
139      guard let data = try? encoder.encode(resp),
140            var str = String(data: data, encoding: .utf8) else { return }
141      str += "\n"
142      if let bytes = str.data(using: .utf8) {
143          output.write(bytes)
144      }
145  }