aoe.c

WinAoE is an open source GPLv3 driver for using AoE (ATA over Ethernet) on Microsoft Windows

/*
  Copyright 2006-2008, V.
  For contact information, see http://winaoe.org/

  This file is part of WinAoE.

  WinAoE is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  WinAoE is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with WinAoE.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "portable.h"
#include <ntddk.h>
#include "aoe.h"
#include "driver.h"
#include "protocol.h"

NTSTATUS STDCALL ZwWaitForSingleObject(IN HANDLE Handle, IN BOOLEAN Alertable, IN PLARGE_INTEGER Timeout OPTIONAL);

// from bus.c
VOID STDCALL BusAddTarget(IN PUCHAR ClientMac, IN PUCHAR ServerMac, USHORT Major, UCHAR Minor, LONGLONG LBASize);
VOID STDCALL BusCleanupTargetList();

// in this file
VOID STDCALL Thread(IN PVOID StartContext);

#ifdef _MSC_VER
#pragma pack(1)
#endif
typedef enum {RequestType, SearchDriveType} TAGTYPE, *PTAGTYPE;

typedef struct _AOE {
  UCHAR ReservedFlag:2;
  UCHAR ErrorFlag:1;
  UCHAR ResponseFlag:1;
  UCHAR Ver:4;
  UCHAR Error;
  USHORT Major;
  UCHAR Minor;
  UCHAR Command;
  UINT Tag;

  UCHAR WriteAFlag:1;
  UCHAR AsyncAFlag:1;
  UCHAR Reserved1AFlag:2;
  UCHAR DeviceHeadAFlag:1;
  UCHAR Reserved2AFlag:1;
  UCHAR ExtendedAFlag:1;
  UCHAR Reserved3AFlag:1;
  union {
    UCHAR Err;
    UCHAR Feature;
  };
  UCHAR Count;
  union {
    UCHAR Cmd;
    UCHAR Status;
  };

  UCHAR Lba0;
  UCHAR Lba1;
  UCHAR Lba2;
  UCHAR Lba3;
  UCHAR Lba4;
  UCHAR Lba5;
  USHORT Reserved;

  UCHAR Data[];
} __attribute__((__packed__)) AOE, *PAOE;
#ifdef _MSC_VER
#pragma pack()
#endif

typedef struct _REQUEST {
  REQUESTMODE Mode;
  ULONG SectorCount;
  PUCHAR Buffer;
  PIRP Irp;
  ULONG TagCount;
  ULONG TotalTags;
} REQUEST, *PREQUEST;

typedef struct _TAG {
  TAGTYPE Type;
  PDEVICEEXTENSION DeviceExtension;
  PREQUEST Request;
  ULONG Id;
  PAOE PacketData;
  ULONG PacketSize;
  LARGE_INTEGER FirstSendTime;
  LARGE_INTEGER SendTime;
  ULONG BufferOffset;
  ULONG SectorCount;
  struct _TAG *Next;
  struct _TAG *Previous;
} TAG, *PTAG;

typedef struct _DISKSEARCH {
  PDEVICEEXTENSION DeviceExtension;
  PTAG Tag;
  struct _DISKSEARCH *Next;
} DISKSEARCH, *PDISKSEARCH;

BOOLEAN Stop = FALSE;
KSPIN_LOCK SpinLock;
KEVENT ThreadSignalEvent;
PTAG TagList = NULL;
PTAG TagListLast = NULL;
PTAG ProbeTag = NULL;
PDISKSEARCH DiskSearchList = NULL;
LONG OutstandingTags = 0;
HANDLE ThreadHandle;

NTSTATUS STDCALL AoEStart() {
  NTSTATUS Status;
  OBJECT_ATTRIBUTES ObjectAttributes;
  PVOID ThreadObject;

  DbgPrint("AoEStart\n");
  if ((ProbeTag = (PTAG)ExAllocatePool(NonPagedPool, sizeof(TAG))) == NULL) {
    DbgPrint("AddStart ExAllocatePool ProbeTag\n");
    return STATUS_INSUFFICIENT_RESOURCES;
  }
  RtlZeroMemory(ProbeTag, sizeof(TAG));
  ProbeTag->PacketSize = sizeof(AOE);
  if ((ProbeTag->PacketData = (PAOE)ExAllocatePool(NonPagedPool, ProbeTag->PacketSize)) == NULL) {
    DbgPrint("AddStart ExAllocatePool ProbeTag->PacketData\n");
    ExFreePool(ProbeTag);
  }
  ProbeTag->SendTime.QuadPart = 0LL;
  RtlZeroMemory(ProbeTag->PacketData, ProbeTag->PacketSize);
  ProbeTag->PacketData->Ver = AOEPROTOCOLVER;
  ProbeTag->PacketData->Major = htons((USHORT)-1);
  ProbeTag->PacketData->Minor = (UCHAR)-1;
  ProbeTag->PacketData->Cmd = 0xec;             // IDENTIFY DEVICE
  ProbeTag->PacketData->Count = 1;

  KeInitializeSpinLock(&SpinLock);
  KeInitializeEvent(&ThreadSignalEvent, SynchronizationEvent, FALSE);

  InitializeObjectAttributes(&ObjectAttributes, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
  if (!NT_SUCCESS(Status = PsCreateSystemThread(&ThreadHandle, THREAD_ALL_ACCESS, &ObjectAttributes, NULL, NULL, Thread, NULL))) return Error("PsCreateSystemThread", Status);
  if (!NT_SUCCESS(Status = ObReferenceObjectByHandle(ThreadHandle, THREAD_ALL_ACCESS, NULL, KernelMode, &ThreadObject, NULL))) {
    ZwClose(ThreadHandle);
    Error("ObReferenceObjectByHandle", Status);
    Stop = TRUE;
    KeSetEvent(&ThreadSignalEvent, 0, FALSE);
  }
  return Status;
}

VOID STDCALL AoEStop() {
  NTSTATUS Status;
  PDISKSEARCH DiskSearch, PreviousDiskSearch;
  PTAG Tag;
  KIRQL Irql;

  DbgPrint("AoEStop\n");
  if (!Stop) {
    Stop = TRUE;
    KeSetEvent(&ThreadSignalEvent, 0, FALSE);
    if (!NT_SUCCESS(Status = ZwWaitForSingleObject(ThreadHandle, FALSE, NULL))) Error("AoEStop ZwWaitForSingleObject", Status);
    ZwClose(ThreadHandle);
  }

  KeAcquireSpinLock(&SpinLock, &Irql);

  DiskSearch = DiskSearchList;
  while (DiskSearch != NULL) {
    KeSetEvent(&DiskSearch->DeviceExtension->Disk.SearchEvent, 0, FALSE);
    PreviousDiskSearch = DiskSearch;
    DiskSearch = DiskSearch->Next;
    ExFreePool(PreviousDiskSearch);
  }

  Tag = TagList;
  while (Tag != NULL) {
    if (Tag->Request != NULL && --Tag->Request->TagCount == 0) {
      Tag->Request->Irp->IoStatus.Information = 0;
      Tag->Request->Irp->IoStatus.Status = STATUS_CANCELLED;
      IoCompleteRequest(Tag->Request->Irp, IO_NO_INCREMENT);
      ExFreePool(Tag->Request);
    }
    if (Tag->Next == NULL) {
      ExFreePool(Tag->PacketData);
      ExFreePool(Tag);
      Tag = NULL;
    } else {
      Tag = Tag->Next;
      ExFreePool(Tag->Previous->PacketData);
      ExFreePool(Tag->Previous);
    }
  }
  TagList = NULL;
  TagListLast = NULL;
  ExFreePool(ProbeTag->PacketData);
  ExFreePool(ProbeTag);

  KeReleaseSpinLock(&SpinLock, Irql);
}

BOOLEAN STDCALL AoESearchDrive(IN PDEVICEEXTENSION DeviceExtension) {
  PDISKSEARCH DiskSearch, DiskSearchWalker, PreviousDiskSearch;
  LARGE_INTEGER Timeout, CurrentTime;
  PTAG Tag, TagWalker;
  KIRQL Irql, InnerIrql;
  LARGE_INTEGER MaxSectorsPerPacketSendTime;
  ULONG MTU;

  if ((DiskSearch = (PDISKSEARCH)ExAllocatePool(NonPagedPool, sizeof(DISKSEARCH))) == NULL) {
    DbgPrint("AoESearchBootDrive ExAllocatePool DiskSearch\n");
    return FALSE;
  }

  DiskSearch->DeviceExtension = DeviceExtension;
  DiskSearch->Next = NULL;
  DeviceExtension->Disk.SearchState = SearchNIC;
  KeResetEvent(&DeviceExtension->Disk.SearchEvent);

  KeAcquireSpinLock(&SpinLock, &Irql);
  if (DiskSearchList == NULL) {
    DiskSearchList = DiskSearch;
  } else {
    DiskSearchWalker = DiskSearchList;
    while (DiskSearchWalker->Next) DiskSearchWalker = DiskSearchWalker->Next;
    DiskSearchWalker->Next = DiskSearch;
  }
  KeReleaseSpinLock(&SpinLock, Irql);

  while (TRUE) {
    Timeout.QuadPart = -500000LL;               // 500.000 * 100ns = 50.000.000 ns = 50ms
    KeWaitForSingleObject(&DeviceExtension->Disk.SearchEvent, Executive, KernelMode, FALSE, &Timeout);
    if (Stop) {
      DbgPrint("AoESearchBootDrive cancled\n");
      return FALSE;
    }

    KeAcquireSpinLock(&DeviceExtension->Disk.SpinLock, &Irql);
    if (DeviceExtension->Disk.SearchState == SearchNIC) {
      if (!ProtocolSearchNIC(DeviceExtension->Disk.ClientMac)) {
        KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
        continue;
      } else {
        DeviceExtension->Disk.MTU = ProtocolGetMTU(DeviceExtension->Disk.ClientMac);
        DeviceExtension->Disk.SearchState = GetSize;
      }
    }

    if (DeviceExtension->Disk.SearchState == GettingSize) {
      KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
      continue;
    }
    if (DeviceExtension->Disk.SearchState == GettingGeometry) {
      KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
      continue;
    }
    if (DeviceExtension->Disk.SearchState == GettingMaxSectorsPerPacket) {
      KeQuerySystemTime(&CurrentTime);
      if (CurrentTime.QuadPart > MaxSectorsPerPacketSendTime.QuadPart + 2500000LL) {    // 2.500.000 * 100ns = 250.000.000 ns = 250ms
        DbgPrint("No reply after 250ms for MaxSectorsPerPacket %d, giving up\n", DeviceExtension->Disk.MaxSectorsPerPacket);
        DeviceExtension->Disk.MaxSectorsPerPacket--;
        DeviceExtension->Disk.SearchState = Done;
      } else {
        KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
        continue;
      }
    }

    if (DeviceExtension->Disk.SearchState == Done) {
      KeAcquireSpinLock(&SpinLock, &InnerIrql);
      TagWalker = TagList;
      while (TagWalker != NULL && TagWalker != Tag) TagWalker = TagWalker->Next;
      if (TagWalker != NULL) {
        if (Tag->Previous == NULL) TagList = Tag->Next;
        else Tag->Previous->Next = Tag->Next;
        if (Tag->Next == NULL) TagListLast = Tag->Previous;
        else Tag->Next->Previous = Tag->Previous;
        OutstandingTags--;
        if (OutstandingTags < 0) DbgPrint("!!OutstandingTags < 0 (AoESearchDrive)!!\n");
        ExFreePool(Tag->PacketData);
        ExFreePool(Tag);
      }

      if (DiskSearchList == NULL) {
        DbgPrint("!!DiskSearchList == NULL!!\n");
      } else {
        DiskSearchWalker = DiskSearchList;
        while (DiskSearchWalker && DiskSearchWalker->DeviceExtension != DeviceExtension) {
          PreviousDiskSearch = DiskSearchWalker;
          DiskSearchWalker = DiskSearchWalker->Next;
        }
        if (DiskSearchWalker) {
          if (DiskSearchWalker == DiskSearchList) {
            DiskSearchList = DiskSearchWalker->Next;
          } else {
            PreviousDiskSearch->Next = DiskSearchWalker->Next;
          }
          ExFreePool(DiskSearchWalker);
        } else {
          DbgPrint("!!Disk not found in DiskSearchList!!\n");
        }
      }

      KeReleaseSpinLock(&SpinLock, InnerIrql);
      KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
        DbgPrint("Disk size: %I64uM cylinders: %I64u heads: %u sectors: %u sectors per packet: %u\n",
        DeviceExtension->Disk.LBADiskSize / 2048,
        DeviceExtension->Disk.Cylinders,
        DeviceExtension->Disk.Heads,
        DeviceExtension->Disk.Sectors,
        DeviceExtension->Disk.MaxSectorsPerPacket
      );
      return TRUE;
    }

    if ((Tag = (PTAG)ExAllocatePool(NonPagedPool, sizeof(TAG))) == NULL) {
      DbgPrint("AoESearchBootDrive ExAllocatePool Tag\n");
      KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
      continue;
    }
    RtlZeroMemory(Tag, sizeof(TAG));
    Tag->Type = SearchDriveType;
    Tag->DeviceExtension = DeviceExtension;
    Tag->PacketSize = sizeof(AOE);
    if ((Tag->PacketData = (PAOE)ExAllocatePool(NonPagedPool, Tag->PacketSize)) == NULL) {
      DbgPrint("AoESearchBootDrive ExAllocatePool Tag->PacketData\n");
      ExFreePool(Tag);
      Tag = NULL;
      KeReleaseSpinLock(&DeviceExtension->Disk.SpinLock, Irql);
      continue;
    }
    RtlZeroMemory(Tag->PacketData, Tag->PacketSize);
    Tag->PacketData->Ver = AOEPROTOCOLVER;
    Tag->PacketData->Major = htons((USHORT)DeviceExtension->Disk.Major);
    Tag->PacketData->Minor = (UCHAR)DeviceExtension->Disk.Minor;
    Tag->PacketData->ExtendedAFlag = TRUE;

    switch (DeviceExtension->Disk.SearchState) {
      case GetSize:
        Tag->PacketData->Cmd = 0xec;               // IDENTIFY DEVICE
        Tag->PacketData->Count = 1;
        DeviceExtension->Disk.SearchState = GettingSize;
        break;
      case GetGeometry:
        Tag->PacketData->Cmd = 0x24;               // READ SECTOR
        Tag->PacketData->Count = 1;
        DeviceExtension->Disk.SearchState = GettingGeometry;
        break;
      case GetMaxSectorsPerPacket:
        Tag->PacketData->Cmd = 0x24;               // READ SECTOR
        Tag->PacketData->Count = (UCHAR)(++DeviceExtension->Disk.MaxSectorsPerPacket);