Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle concurrent requests to socket #495

Merged
merged 4 commits into from
Dec 11, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions Sources/Packages/Sources/SecretAgentKit/Agent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ public class Agent {
private let writer = OpenSSHKeyWriter()
private let requestTracer = SigningRequestTracer()
private let certificateHandler = OpenSSHCertificateHandler()
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent.agent", category: "")
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "Agent")

/// Initializes an agent with a store list and a witness.
/// - Parameters:
Expand All @@ -35,7 +35,7 @@ extension Agent {
/// - writer: A ``FileHandleWriter`` to write the response to.
/// - Return value:
/// - Boolean if data could be read
@discardableResult public func handle(reader: FileHandleReader, writer: FileHandleWriter) -> Bool {
@discardableResult @Sendable public func handle(reader: FileHandleReader, writer: FileHandleWriter) async -> Bool {
logger.debug("Agent handling new data")
let data = Data(reader.availableData)
guard data.count > 4 else { return false}
Expand All @@ -47,12 +47,12 @@ extension Agent {
}
logger.debug("Agent handling request of type \(requestType.debugDescription)")
let subData = Data(data[5...])
let response = handle(requestType: requestType, data: subData, reader: reader)
let response = await handle(requestType: requestType, data: subData, reader: reader)
writer.write(response)
return true
}

func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) -> Data {
func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) async -> Data {
// Depending on the launch context (such as after macOS update), the agent may need to reload secrets before acting
reloadSecretsIfNeccessary()
var response = Data()
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import Foundation

/// Protocol abstraction of the reading aspects of FileHandle.
public protocol FileHandleReader {
public protocol FileHandleReader: Sendable {

/// Gets data that is available for reading.
var availableData: Data { get }
Expand All @@ -13,7 +13,7 @@ public protocol FileHandleReader {
}

/// Protocol abstraction of the writing aspects of FileHandle.
public protocol FileHandleWriter {
public protocol FileHandleWriter: Sendable {

/// Writes data to the handle.
func write(_ data: Data)
Expand Down
56 changes: 39 additions & 17 deletions Sources/Packages/Sources/SecretAgentKit/SocketController.swift
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,24 @@ public class SocketController {
private var port: SocketPort?
/// A handler that will be notified when a new read/write handle is available.
/// False if no data could be read
public var handler: ((FileHandleReader, FileHandleWriter) -> Bool)?
public var handler: ((FileHandleReader, FileHandleWriter) async -> Bool)?
/// Logger.
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "SocketController")


/// Initializes a socket controller with a specified path.
/// - Parameter path: The path to use as a socket.
public init(path: String) {
Logger().debug("Socket controller setting up at \(path)")
logger.debug("Socket controller setting up at \(path)")
if let _ = try? FileManager.default.removeItem(atPath: path) {
Logger().debug("Socket controller removed existing socket")
logger.debug("Socket controller removed existing socket")
}
let exists = FileManager.default.fileExists(atPath: path)
assert(!exists)
Logger().debug("Socket controller path is clear")
logger.debug("Socket controller path is clear")
port = socketPort(at: path)
configureSocket(at: path)
Logger().debug("Socket listening at \(path)")
logger.debug("Socket listening at \(path)")
}

/// Configures the socket and a corresponding FileHandle.
Expand All @@ -35,7 +37,7 @@ public class SocketController {
fileHandle = FileHandle(fileDescriptor: port.socket, closeOnDealloc: true)
NotificationCenter.default.addObserver(self, selector: #selector(handleConnectionAccept(notification:)), name: .NSFileHandleConnectionAccepted, object: nil)
NotificationCenter.default.addObserver(self, selector: #selector(handleConnectionDataAvailable(notification:)), name: .NSFileHandleDataAvailable, object: nil)
fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!])
fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.Mode.common])
}

/// Creates a SocketPort for a path.
Expand Down Expand Up @@ -65,25 +67,45 @@ public class SocketController {
/// Handles a new connection being accepted, invokes the handler, and prepares to accept new connections.
/// - Parameter notification: A `Notification` that triggered the call.
@objc func handleConnectionAccept(notification: Notification) {
Logger().debug("Socket controller accepted connection")
logger.debug("Socket controller accepted connection")
guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return }
_ = handler?(new, new)
new.waitForDataInBackgroundAndNotify()
fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!])
Task {
_ = await handler?(new, new)
await new.waitForDataInBackgroundAndNotifyOnMainActor()
await fileHandle?.acceptConnectionInBackgroundAndNotifyOnMainActor()
}
}

/// Handles a new connection providing data and invokes the handler callback.
/// - Parameter notification: A `Notification` that triggered the call.
@objc func handleConnectionDataAvailable(notification: Notification) {
Logger().debug("Socket controller has new data available")
logger.debug("Socket controller has new data available")
guard let new = notification.object as? FileHandle else { return }
Logger().debug("Socket controller received new file handle")
if((handler?(new, new)) == true) {
Logger().debug("Socket controller handled data, wait for more data")
new.waitForDataInBackgroundAndNotify()
} else {
Logger().debug("Socket controller called with empty data, socked closed")
logger.debug("Socket controller received new file handle")
Task {
if((await handler?(new, new)) == true) {
logger.debug("Socket controller handled data, wait for more data")
await new.waitForDataInBackgroundAndNotifyOnMainActor()
} else {
logger.debug("Socket controller called with empty data, socked closed")
}
}
}

}

extension FileHandle {

/// Ensures waitForDataInBackgroundAndNotify will be called on the main actor.
@MainActor func waitForDataInBackgroundAndNotifyOnMainActor() {
waitForDataInBackgroundAndNotify()
}


/// Ensures acceptConnectionInBackgroundAndNotify will be called on the main actor.
/// - Parameter modes: the runloop modes to use.
@MainActor func acceptConnectionInBackgroundAndNotifyOnMainActor(forModes modes: [RunLoop.Mode]? = [RunLoop.Mode.common]) {
acceptConnectionInBackgroundAndNotify(forModes: modes)
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import OSLog
public class OpenSSHCertificateHandler {

private let publicKeyFileStoreController = PublicKeyFileStoreController(homeDirectory: NSHomeDirectory())
private let logger = Logger()
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "OpenSSHCertificateHandler")
private let writer = OpenSSHKeyWriter()
private var keyBlobsAndNames: [AnySecret: (Data, Data)] = [:]

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import OSLog
/// Controller responsible for writing public keys to disk, so that they're easily accessible by scripts.
public class PublicKeyFileStoreController {

private let logger = Logger()
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "PublicKeyFileStoreController")
private let directory: String
private let keyWriter = OpenSSHKeyWriter()

Expand Down
36 changes: 18 additions & 18 deletions Sources/Packages/Tests/SecretAgentKitTests/AgentTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -10,39 +10,39 @@ class AgentTests: XCTestCase {

// MARK: Identity Listing

func testEmptyStores() {
func testEmptyStores() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestIdentities)
let agent = Agent(storeList: SecretStoreList())
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(stubWriter.data, Constants.Responses.requestIdentitiesEmpty)
}

func testIdentitiesList() {
func testIdentitiesList() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestIdentities)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
let agent = Agent(storeList: list)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(stubWriter.data, Constants.Responses.requestIdentitiesMultiple)
}

// MARK: Signatures

func testNoMatchingIdentities() {
func testNoMatchingIdentities() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignatureWithNoneMatching)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
let agent = Agent(storeList: list)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
// XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
}

func testSignature() {
func testSignature() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
let requestReader = OpenSSHReader(data: Constants.Requests.requestSignature[5...])
_ = requestReader.readNextChunk()
let dataToSign = requestReader.readNextChunk()
let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
let agent = Agent(storeList: list)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
let outer = OpenSSHReader(data: stubWriter.data[5...])
let payload = outer.readNextChunk()
let inner = OpenSSHReader(data: payload)
Expand Down Expand Up @@ -76,18 +76,18 @@ class AgentTests: XCTestCase {

// MARK: Witness protocol

func testWitnessObjectionStopsRequest() {
func testWitnessObjectionStopsRequest() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
let witness = StubWitness(speakNow: { _,_ in
return true
}, witness: { _, _ in })
let agent = Agent(storeList: list, witness: witness)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
}

func testWitnessSignature() {
func testWitnessSignature() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
var witnessed = false
Expand All @@ -97,11 +97,11 @@ class AgentTests: XCTestCase {
witnessed = true
})
let agent = Agent(storeList: list, witness: witness)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertTrue(witnessed)
}

func testRequestTracing() {
func testRequestTracing() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
var speakNowTrace: SigningRequestProvenance! = nil
Expand All @@ -113,7 +113,7 @@ class AgentTests: XCTestCase {
witnessTrace = trace
})
let agent = Agent(storeList: list, witness: witness)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(witnessTrace, speakNowTrace)
XCTAssertEqual(witnessTrace.origin.displayName, "Finder")
XCTAssertEqual(witnessTrace.origin.validSignature, true)
Expand All @@ -122,22 +122,22 @@ class AgentTests: XCTestCase {

// MARK: Exception Handling

func testSignatureException() {
func testSignatureException() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
let store = list.stores.first?.base as! Stub.Store
store.shouldThrow = true
let agent = Agent(storeList: list)
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
}

// MARK: Unsupported

func testUnhandledAdd() {
func testUnhandledAdd() async {
let stubReader = StubFileHandleReader(availableData: Constants.Requests.addIdentity)
let agent = Agent(storeList: SecretStoreList())
agent.handle(reader: stubReader, writer: stubWriter)
await agent.handle(reader: stubReader, writer: stubWriter)
XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
}

Expand Down
3 changes: 2 additions & 1 deletion Sources/SecretAgent/AppDelegate.swift
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,10 @@ class AppDelegate: NSObject, NSApplicationDelegate {
return SocketController(path: path)
}()
private var updateSink: AnyCancellable?
private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "AppDelegate")

func applicationDidFinishLaunching(_ aNotification: Notification) {
Logger().debug("SecretAgent finished launching")
logger.debug("SecretAgent finished launching")
DispatchQueue.main.async {
self.socketController.handler = self.agent.handle(reader:writer:)
}
Expand Down
10 changes: 6 additions & 4 deletions Sources/Secretive/Controllers/LaunchAgentController.swift
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@ import SecretKit

struct LaunchAgentController {

private let logger = Logger(subsystem: "com.maxgoedjen.secretive", category: "LaunchAgentController")

func install(completion: (() -> Void)? = nil) {
Logger().debug("Installing agent")
logger.debug("Installing agent")
_ = setEnabled(false)
// This is definitely a bit of a "seems to work better" thing but:
// Seems to more reliably hit if these are on separate runloops, otherwise it seems like it sometimes doesn't kill old
Expand All @@ -20,7 +22,7 @@ struct LaunchAgentController {
}

func forceLaunch(completion: ((Bool) -> Void)?) {
Logger().debug("Agent is not running, attempting to force launch")
logger.debug("Agent is not running, attempting to force launch")
let url = Bundle.main.bundleURL.appendingPathComponent("Contents/Library/LoginItems/SecretAgent.app")
let config = NSWorkspace.OpenConfiguration()
config.activates = false
Expand All @@ -29,9 +31,9 @@ struct LaunchAgentController {
completion?(error == nil)
}
if let error = error {
Logger().error("Error force launching \(error.localizedDescription)")
logger.error("Error force launching \(error.localizedDescription)")
} else {
Logger().debug("Agent force launched")
logger.debug("Agent force launched")
}
}
}
Expand Down
Loading