Merge pull request #149 from kishikawakatsumi/freespace

Added function to check free space on shared drives

Sources/SMBClient/Messages/FileInformationClasses/FileInfoClass.swift

@@ -48,6 +48,18 @@ public enum FileInfoClass: UInt8 {
   case fileStreamInformation = 0x16
   case fileTrackingInformation = 0x24
   case fileValidDataLengthInformation = 0x27
+
+  static let fileFsVolumeInformation = fileDirectoryInformation
+  static let fileFsLabelInformation = fileFullDirectoryInformation
+  static let fileFsSizeInformation = fileBothDirectoryInformation
+  static let fileFsDeviceInformation = fileBasicInformation
+  static let fileFsAttributeInformation = fileStandardInformation
+  static let fileFsControlInformation = fileInternalInformation
+  static let fileFsFullSizeInformation = fileEaInformation
+  static let fileFsObjectIdInformation = fileAccessInformation
+  static let fileFsDriverPathInformation = fileNameInformation
+  static let fileFsVolumeFlagsInformation = fileModeInformation
+  static let fileFsSectorSizeInformation: FileInfoClass = fileAlignmentInformation
 }
 
 extension FileInfoClass: CustomDebugStringConvertible {

Sources/SMBClient/Messages/FileSystemInformationClasses/FileFsDeviceInformation.swift

@@ -0,0 +1,39 @@
+import Foundation
+
+public struct FileFsDeviceInformation {
+  public let deviceType: UInt32
+  public let characteristics: Characteristics
+
+  public enum DeviceType: UInt32 {
+    case cdRom = 0x00000002
+    case disk = 0x00000007
+  }
+
+  public struct Characteristics: OptionSet {
+    public let rawValue: UInt32
+
+    public init(rawValue: UInt32) {
+      self.rawValue = rawValue
+    }
+
+    public static let removableMedia = Characteristics(rawValue: 0x00000001)
+    public static let readOnlyDevice = Characteristics(rawValue: 0x00000002)
+    public static let floppyDiskette = Characteristics(rawValue: 0x00000004)
+    public static let writeOnceMedia = Characteristics(rawValue: 0x00000008)
+    public static let remoteDevice = Characteristics(rawValue: 0x00000010)
+    public static let deviceIsMounted = Characteristics(rawValue: 0x00000020)
+    public static let virtualVolume = Characteristics(rawValue: 0x00000040)
+    public static let deviceSecureOpen = Characteristics(rawValue: 0x00000100)
+    public static let characteristicTsDevice = Characteristics(rawValue: 0x00001000)
+    public static let characteristicWebDavDevice = Characteristics(rawValue: 0x00002000)
+    public static let deviceAllowAppContainerTraversal = Characteristics(rawValue: 0x00020000)
+    public static let portableDevice = Characteristics(rawValue: 0x00004000)
+  }
+
+  public init(data: Data) {
+    let reader = ByteReader(data)
+
+    deviceType = reader.read()
+    characteristics = Characteristics(rawValue: reader.read())
+  }
+}

Sources/SMBClient/Messages/FileSystemInformationClasses/FileFsLabelInformation.swift

@@ -0,0 +1,14 @@
+import Foundation
+
+public struct FileFsLabelInformation {
+  public let volumeLabelLength: UInt32
+  public let volumeLabel: String
+
+  public init(data: Data) {
+    let reader = ByteReader(data)
+
+    volumeLabelLength = reader.read()
+    let volumeLabelData = reader.read(count: Int(volumeLabelLength))
+    volumeLabel = String(data: volumeLabelData, encoding: .utf16LittleEndian) ?? volumeLabelData.hex
+  }
+}

Sources/SMBClient/Messages/FileSystemInformationClasses/FileFsSizeInformation.swift

@@ -0,0 +1,17 @@
+import Foundation
+
+public struct FileFsSizeInformation {
+  public let totalAllocationUnits: UInt64
+  public let availableAllocationUnits: UInt64
+  public let sectorsPerAllocationUnit: UInt32
+  public let bytesPerSector: UInt32
+
+  public init(data: Data) {
+    let reader = ByteReader(data)
+
+    totalAllocationUnits = reader.read()
+    availableAllocationUnits = reader.read()
+    sectorsPerAllocationUnit = reader.read()
+    bytesPerSector = reader.read()
+  }
+}

Sources/SMBClient/Messages/FileSystemInformationClasses/FileFsVolumeInformation.swift

@@ -0,0 +1,22 @@
+import Foundation
+
+public struct FileFsVolumeInformation {
+  public let volumeCreationTime: UInt64
+  public let volumeSerialNumber: UInt32
+  public let volumeLabelLength: UInt32
+  public let supportsObjects: Bool
+  public let reserved: UInt8
+  public let volumeLabel: String
+
+  public init(data: Data) {
+    let reader = ByteReader(data)
+
+    volumeCreationTime = reader.read()
+    volumeSerialNumber = reader.read()
+    volumeLabelLength = reader.read()
+    supportsObjects = reader.read() == 1
+    reserved = reader.read()
+    let volumeLabelData = reader.read(count: Int(volumeLabelLength))
+    volumeLabel = String(data: volumeLabelData, encoding: .utf16LittleEndian) ?? volumeLabelData.hex
+  }
+}

Sources/SMBClient/SMBClient.swift

@@ -170,6 +170,20 @@ public class SMBClient {
     FileWriter(session: session, path: Pathname.normalize(path))
   }
 
+  public func availableSpace() async throws -> UInt64 {
+    let response = try await session.queryInfo(path: "", infoType: .fileSystem, fileInfoClass: .fileFsSizeInformation)
+
+    let sizeInformation = FileFsSizeInformation(data: response.buffer)
+    let availableAllocationUnits = sizeInformation.availableAllocationUnits
+    let sectorsPerAllocationUnit = sizeInformation.sectorsPerAllocationUnit
+    let bytesPerSector = sizeInformation.bytesPerSector
+
+    let bytesPerAllocationUnit = UInt64(sectorsPerAllocationUnit * bytesPerSector)
+    let availableSpaceBytes = availableAllocationUnits * bytesPerAllocationUnit
+
+    return availableSpaceBytes
+  }
+
   public func keepAlive() async throws -> Echo.Response {
     try await session.echo()
   }

Sources/SMBClient/Session.swift

@@ -412,7 +412,7 @@ public class Session {
     }
   }
 
-  public func queryInfo(path: String, fileInfoClass: FileInfoClass = .fileAllInformation) async throws -> QueryInfo.Response {
+  public func queryInfo(path: String, infoType: InfoType = .file, fileInfoClass: FileInfoClass = .fileAllInformation) async throws -> QueryInfo.Response {
     let createRequest = Create.Request(
       messageId: messageId.next(),
       treeId: treeId,
@@ -429,7 +429,7 @@ public class Session {
       messageId: messageId.next(),
       treeId: treeId,
       sessionId: sessionId,
-      infoType: .file,
+      infoType: infoType,
       fileInfoClass: fileInfoClass,
       fileId: temporaryUUID
     )

Tests/SMBClientTests/SessionTests.swift

@@ -80,6 +80,22 @@ final class SessionTests: XCTestCase {
     try await session.logoff()
   }
 
+  func testQueryInfo01() async throws {
+    let session = Session(host: "localhost", port: 4445)
+
+    try await session.connect()
+    try await session.negotiate()
+
+    try await session.sessionSetup(username: "alice", password: "alipass")
+    try await session.treeConnect(path: "Alice Share")
+
+    let response = try await session.queryInfo(path: "", infoType: .fileSystem, fileInfoClass: .fileFsSizeInformation)
+    XCTAssertTrue(NTStatus(response.header.status) == .success)
+
+    try await session.treeDisconnect()
+    try await session.logoff()
+  }
+
   func testEcho() async throws {
     let session = Session(host: "localhost", port: 4445)