From bdf049d9e433e05c35ddc01f177e38254ac449aa Mon Sep 17 00:00:00 2001
From: Colin <colin@uninsane.org>
Date: Thu, 28 Sep 2023 20:55:18 +0000
Subject: [PATCH] moby: wowlan: also wake on ARP requests (experimental)

---
 hosts/by-name/moby/default.nix |  1 +
 modules/wowlan.nix             | 76 +++++++++++++++++++++++++++-------
 2 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/hosts/by-name/moby/default.nix b/hosts/by-name/moby/default.nix
index 875c9098..1270db9f 100644
--- a/hosts/by-name/moby/default.nix
+++ b/hosts/by-name/moby/default.nix
@@ -27,6 +27,7 @@
   sane.wowlan.patterns = [
     { ipv4.destPort = 22; }  # wake on SSH
     { ipv4.srcPort = 2587; }  # wake on `ntfy-sh` push from servo
+    { arp.queryIp = [ 10 78 79 54 ]; }  # wake when somebody is doing an ARP query against us
   ];
 
   # XXX colin: phosh doesn't work well with passwordless login,
diff --git a/modules/wowlan.nix b/modules/wowlan.nix
index b742a654..79fb3656 100644
--- a/modules/wowlan.nix
+++ b/modules/wowlan.nix
@@ -43,6 +43,14 @@ let
           (e.g. sshing *into* this machine).
         '';
       };
+      arp.queryIp = mkOption {
+        type = types.nullOr (types.listOf types.int);
+        default = null;
+        description = ''
+          IP address being queried.
+          e.g. `[ 192 168 0 100 ]`
+        '';
+      };
     };
   };
   # bytesToStr: [ u8|null ] -> String
@@ -59,6 +67,46 @@ let
     "0" + (lib.toHexString b)
   else
     lib.toHexString b;
+
+  etherTypes.ipv4 = [ 08 00 ];  # 0x0800 = IPv4
+  etherTypes.arp = [ 08 06 ];  # 0x0806 = ARP
+  formatEthernetFrame = ethertype: dataBytes: let
+    bytes = [
+      # ethernet frame: <https://en.wikipedia.org/wiki/Ethernet_frame#Structure>
+      ## dest MAC address (this should be the device's MAC, but i think that's implied?)
+      null null null null null null
+      ## src MAC address
+      null null null null null null
+      ## ethertype: <https://en.wikipedia.org/wiki/EtherType#Values>
+    ] ++ etherTypes."${ethertype}"
+      ++ dataBytes;
+  in bytesToStr bytes;
+
+  formatArpPattern = pat:
+    formatEthernetFrame "arp" ([
+      # ARP frame: <https://en.wikipedia.org/wiki/Address_Resolution_Protocol#Packet_structure>
+      ## hardware type
+      null null
+      ## protocol type. same coding as EtherType
+      08 00  # 0x0800 = IPv4
+      ## hardware address length (i.e. MAC)
+      06
+      ## protocol address length (i.e. IP address)
+      04
+      ## operation
+      00 01  # 0x0001 = request
+      ## sender hardware address
+      null null null null null null
+      ## sender protocol address
+      null null null null
+      ## target hardware address
+      ## this is left as "Don't Care" because the packets we want to match
+      ## are those mapping protocol addr -> hw addr.
+      ## sometimes clients do include this field if they've seen the address before though
+      null null null null null null
+      ## target protocol address
+    ] ++ pat.queryIp);
+
   # formatIpv4Pattern: patternOpts.ipv4 -> String
   # produces a string like this (example matches source port=0x0a1b):
   # "-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:-:0a:1b:-:-"
@@ -75,16 +123,8 @@ let
     } else {
       high = null; low = null;
     };
-
-    bytes = [
-      # ethernet frame: <https://en.wikipedia.org/wiki/Ethernet_frame#Structure>
-      ## dest MAC address (this should be the device's MAC, but i think that's implied?)
-      null null null null null null
-      ## src MAC address
-      null null null null null null
-      ## ethertype: <https://en.wikipedia.org/wiki/EtherType#Values>
-      08 00  # 0x0800 = IPv4
-
+  in
+    formatEthernetFrame "ipv4" [
       # IP frame: <https://en.wikipedia.org/wiki/Internet_Protocol_version_4#Header>
       ## Version, Internet Header Length. 0x45 = 69 decimal
       null  # should be 69 (0x45), but fails to wake if i include this
@@ -112,7 +152,6 @@ let
       destPortBytes.high destPortBytes.low
       ## rest is Don't Care
     ];
-  in bytesToStr bytes;
 in
 {
   options = with lib; {
@@ -136,17 +175,22 @@ in
       description = "configure the WiFi chip to wake the system on specific activity";
       path = [ pkgs.iw pkgs.wirelesstools ];
       script = let
-        writePattern = pat: ''
-          iwpriv wlan0 wow_set_pattern pattern=${formatIpv4Pattern pat.ipv4}
+        extractPatterns = pat: lib.flatten [
+          (lib.optional (pat.ipv4 != { destPort = null; srcPort = null; }) (formatIpv4Pattern pat.ipv4))
+          (lib.optional (pat.arp != { queryIp = null; }) (formatArpPattern pat.arp))
+        ];
+        allPatterns = lib.flatten (builtins.map extractPatterns cfg.patterns);
+        encodePattern = pat: ''
+          iwpriv wlan0 wow_set_pattern pattern=${pat}
         '';
-        writePatterns = lib.concatStringsSep
+        encodedPatterns = lib.concatStringsSep
           "\n"
-          (builtins.map writePattern cfg.patterns);
+          (builtins.map encodePattern allPatterns);
       in ''
         set -x
         iw phy0 wowlan enable any
         iwpriv wlan0 wow_set_pattern clean
-        ${writePatterns}
+        ${encodedPatterns}
       '';
       serviceConfig = {
         # TODO: re-run this periodically, just to be sure?