Merge pull request #20 from ntoskrnl7/features/ntl

Features/ntl

include/.internal/version

@@ -26,6 +26,6 @@ Environment:
 
 #define CRTSYS_VERSION_MAJOR 0
 #define CRTSYS_VERSION_MINOR 1
-#define CRTSYS_VERSION_PATCH 6
+#define CRTSYS_VERSION_PATCH 7
 
 #endif // _CRTSYS_VERSION_

include/ntl/device

@@ -0,0 +1,146 @@
+/**
+ * @file device
+ * @author jungkwang.lee ([email protected])
+ * @brief
+ *
+ * @copyright Copyright (c) 2022 NT Template Library Authoers.
+ *
+ */
+#pragma once
+
+#include "except"
+#include <functional>
+#include <stdexcept>
+#include <string>
+
+#ifndef _NTDDK_
+#include <wdm.h>
+#endif
+
+namespace ntl {
+class device {
+  friend class device_dispatch_invoker;
+
+public:
+  using on_device_control_fn = std::function<void(uint32_t, const uint8_t *,
+                                                  size_t, uint8_t *, size_t *)>;
+
+private:
+  friend class driver;
+
+  struct context_base {
+    context_base(const std::type_info &type) : type(type) {}
+    const std::type_info &type;
+    on_device_control_fn on_device_control;
+  };
+
+  template <typename T> class context;
+
+  template <> struct context<void> : public context_base {
+    context() : context_base(typeid(void)) {}
+  };
+
+  template <typename T> struct context : public context_base {
+    context() : context_base(typeid(T)) {}
+    T data;
+  };
+
+public:
+  using type_t = DEVICE_TYPE;
+
+  class options {
+    friend class driver;
+    friend class device;
+
+  public:
+    options() : type_(FILE_DEVICE_UNKNOWN), exclusive_(false) {}
+
+    options &name(const std::wstring &name) {
+      name_ = name;
+      return *this;
+    }
+    options &type(type_t type) {
+      type_ = type;
+      return *this;
+    }
+    options &exclusive(bool exclusive = true) {
+      exclusive_ = exclusive;
+      return *this;
+    }
+
+  private:
+    bool exclusive_;
+    std::wstring name_;
+    type_t type_;
+  };
+
+protected:
+  device(PDEVICE_OBJECT device, const options &opts,
+         std::function<void()> finalizer)
+      : object_(device), name_(opts.name_), finalizer_(finalizer) {}
+
+private:
+  device(const device &) = delete;
+
+public:
+  device() : object_(nullptr) {}
+
+  device(device &&other) { *this = std::move(other); }
+
+  device &operator=(device &&rhs) {
+    object_ = rhs.detach();
+    name_.swap(rhs.name_);
+    finalizer_ = std::move(rhs.finalizer_);
+    return *this;
+  }
+
+  ~device() {
+    auto obj = detach();
+    if (obj) {
+      if (finalizer_)
+        finalizer_();
+      IoDeleteDevice(obj);
+    }
+  }
+
+private:
+  template <typename Extension> context<Extension> *get_context() {
+    return reinterpret_cast<context<Extension> *>(object_->DeviceExtension);
+  }
+
+public:
+  template <typename Extension> Extension &extension() {
+    auto ext = get_context<Extension>();
+    if (ext->type != typeid(Extension))
+      throw std::runtime_error("bad type");
+    return ext->data;
+  }
+
+  device &on_device_control(on_device_control_fn &&f) {
+    get_context<void>()->on_device_control = f;
+    return *this;
+  }
+
+  const std::wstring &name() const {
+    if (!object_)
+      throw std::runtime_error("invalid device object.");
+    return name_;
+  }
+
+  type_t type() const {
+    if (!object_)
+      throw std::runtime_error("invalid device object.");
+    return object_->DeviceType;
+  }
+
+  PDEVICE_OBJECT detach() {
+    return reinterpret_cast<PDEVICE_OBJECT>(
+        InterlockedExchangePointer(reinterpret_cast<PVOID *>(&object_), NULL));
+  }
+
+private:
+  PDEVICE_OBJECT object_;
+  std::wstring name_;
+  std::function<void()> finalizer_;
+};
+} // namespace ntl

include/ntl/driver

@@ -11,128 +11,109 @@
 #ifndef _NTDDK_
 #include <wdm.h>
 #endif
+#ifndef ClearFlag
+#define ClearFlag(_F, _SF) ((_F) &= ~(_SF))
+#endif
+
+#include "device"
+#include "except"
 #include "status"
+#include "unicode_string"
+#include <algorithm>
 #include <functional>
-#include <string>
+#include <memory>
+#include <unordered_map>
 
 namespace ntl {
-class unicode_string {
-public:
-  unicode_string(std::wstring &&str) : str_(std::move(str)) {
-    value_.Buffer = &str_[0];
-    value_.MaximumLength = value_.Length = (USHORT)str_.size() * sizeof(WCHAR);
-  }
-  unicode_string(const std::wstring &str) : str_(str) {
-    value_.Buffer = &str_[0];
-    value_.MaximumLength = value_.Length = (USHORT)str_.size() * sizeof(WCHAR);
-  }
-  unicode_string(std::wstring &str) {
-    value_.Buffer = &str[0];
-    value_.MaximumLength = value_.Length = (USHORT)str.size() * sizeof(WCHAR);
-  }
-
-  UNICODE_STRING &operator*() { return value_; }
-
-private:
-  std::wstring str_;
-  UNICODE_STRING value_;
-};
-
-class device {
-  friend class driver;
-
-public:
-  using type = ULONG;
-
-protected:
-  device(PDEVICE_OBJECT device) : object_(device) {}
-
-private:
-  device(const device &) = delete;
-
-public:
-  device() : object_(NULL) {}
-
-  device(device &&other) { *this = std::move(other); }
-
-  device &operator=(device &&rhs) {
-    object_ = rhs.object_;
-    rhs.object_ = NULL;
-    return *this;
-  }
-
-  ~device() {
-    if (object_)
-      IoDeleteDevice(object_);
-  }
-
-public:
-  PDEVICE_OBJECT detach() {
-    return (PDEVICE_OBJECT)InterlockedExchangePointer((PVOID *)&object_, NULL);
+namespace detail {
+namespace driver {
+inline ntl::except make_exception(NTSTATUS status) {
+  const char *message;
+  switch (status) {
+  case STATUS_INSUFFICIENT_RESOURCES:
+    message = " Insufficient system resources exist to complete the API.";
+    break;
+  case STATUS_OBJECT_NAME_COLLISION:
+    message = "Object Name already exists.";
+    break;
+  default:
+    message = "Unknown error.";
+    break;
   }
-
-#ifdef NTL_EXPORT_OBJECT_PTR
-public:
-#else
-protected:
-#endif
-  PDEVICE_OBJECT ptr() const { return object_; }
-
-private:
-  PDEVICE_OBJECT object_;
-};
+  return std::move(ntl::except(status, message));
+}
+} // namespace driver
+} // namespace detail
 
 class driver {
-public:
+  friend class driver_initializer;
   friend class driver_unload_invoker;
-  friend class device_dispatch_invoker;
-
-  using unload_routine = std::function<void()>;
-
-  using devcie_control_routine =
-      std::function<void(uint32_t, const uint8_t *, size_t, uint8_t *, size_t *)>;
 
+private:
   driver(PDRIVER_OBJECT driver)
       : object_(driver), name_(driver->DriverName.Buffer,
                                driver->DriverName.Length / sizeof(WCHAR)) {}
 
-  device create_device(const std::wstring &name, device::type type) {
-    unicode_string dev_name(L"\\Device\\" + name);
-    PDEVICE_OBJECT dev = NULL;
-    status s = IoCreateDevice(object_, 0, &*dev_name, type, 0, FALSE, &dev);
-    return dev;
-  }
+public:
+  using unload_routine = std::function<void()>;
+
+  using devcie_control_routine = std::function<void(
+      uint32_t, const uint8_t *, size_t, uint8_t *, size_t *)>;
 
-  void on_device_control(devcie_control_routine &&f) {
-    devcie_control_routine_ = f;
+  template <typename Extension>
+  std::shared_ptr<device> create_device(device::options &opts) {
+    PDEVICE_OBJECT dev = NULL;
+    status s = [&dev, &opts, this]() {
+      ULONG extension_size = (ULONG)sizeof(device::context<Extension>);
+      if (opts.name_.empty()) {
+        return IoCreateDevice(object_, extension_size, NULL, opts.type_, 0,
+                              opts.exclusive_, &dev);
+      } else {
+        unicode_string dev_name(L"\\Device\\" + opts.name_);
+        return IoCreateDevice(object_, extension_size, &*dev_name, opts.type_,
+                              0, opts.exclusive_, &dev);
+      }
+    }();
+    if (s.is_err())
+      throw detail::driver::make_exception(s);
+
+    new (dev->DeviceExtension) device::context<Extension>();
+
+    auto ptr = std::make_shared<device>(std::move(device(dev, opts, [dev]() {
+      reinterpret_cast<device::context<Extension> *>(dev->DeviceExtension)
+          ->~context();
+    })));
+    devices_.insert({dev, ptr});
+
+    ClearFlag(dev->Flags, DO_DEVICE_INITIALIZING);
+
+    return ptr;
   }
 
   void on_unload(unload_routine &&f) { unload_routine_ = f; }
 
-  PDRIVER_OBJECT detach() {
-    return (PDRIVER_OBJECT)InterlockedExchangePointer((PVOID *)&object_, NULL);
+  std::shared_ptr<device> devices(PDEVICE_OBJECT ptr) noexcept {
+    try {
+      return devices_[ptr];
+    } catch (...) {
+      return nullptr;
+    }
+  };
+
+  std::vector<std::shared_ptr<device>> devices() {
+    std::vector<std::shared_ptr<device>> vec;
+    std::transform(devices_.begin(), devices_.end(), vec.end(),
+                   [](auto &item) { return item.second; });
+    return vec;
   }
 
   const std::wstring &name() const { return name_; }
 
-#ifdef NTL_EXPORT_OBJECT_PTR
-public:
-#else
-protected:
-#endif
-  PDRIVER_OBJECT ptr() const { return object_; }
-
-private:
-  void unload() {
-    if (unload_routine_)
-      unload_routine_();
-  }
-
 private:
   PDRIVER_OBJECT object_;
   std::wstring name_;
 
-  devcie_control_routine devcie_control_routine_;
+  std::unordered_map<PDEVICE_OBJECT, std::shared_ptr<device>> devices_;
   unload_routine unload_routine_;
 };
 

include/ntl/expand_stack

@@ -115,10 +115,11 @@ inline ntl::except make_exception(NTSTATUS status) {
   const char *message;
   switch (status) {
   case STATUS_INVALID_PARAMETER_3:
-    message = "The stack_size parameter is greater than MAXIMUM_EXPANSION_SIZE";
+    message =
+        "The stack_size parameter is greater than MAXIMUM_EXPANSION_SIZE.";
     break;
   case STATUS_INVALID_PARAMETER_4:
-    message = "The routine was called at IRQL = DISPATCH_LEVEL";
+    message = "The routine was called at IRQL = DISPATCH_LEVEL.";
     break;
   case STATUS_NO_MEMORY:
     message = "Not enough memory is available to expand the stack.";
@@ -128,7 +129,7 @@ inline ntl::except make_exception(NTSTATUS status) {
               "system's internal limits on stack space.";
     break;
   default:
-    message = "Unknown error";
+    message = "Unknown error.";
     break;
   }
   return std::move(ntl::except(status, message));