feat: improve wloc response compatibility

2026-09-24 23:51:58 +08:00 · 2026-08-09 16:23:41 +08:00
parent 7cf38363cf
commit b0d2f0e8e4
4 changed files with 235 additions and 7 deletions
@@ -44,6 +44,29 @@ func testFrame(payload []byte) []byte {
 	return out
 }

+func testARPCFrame(payload, suffix []byte) ([]byte, int) {
+	var out []byte
+	var version [2]byte
+	binary.BigEndian.PutUint16(version[:], 1)
+	out = append(out, version[:]...)
+	for _, value := range [][]byte{[]byte("zh_CN"), []byte("com.apple.locationd"), []byte("20A123")} {
+		var length [2]byte
+		binary.BigEndian.PutUint16(length[:], uint16(len(value)))
+		out = append(out, length[:]...)
+		out = append(out, value...)
+	}
+	var functionID [4]byte
+	binary.BigEndian.PutUint32(functionID[:], 1)
+	out = append(out, functionID[:]...)
+	lengthOffset := len(out)
+	var payloadLength [4]byte
+	binary.BigEndian.PutUint32(payloadLength[:], uint32(len(payload)))
+	out = append(out, payloadLength[:]...)
+	out = append(out, payload...)
+	out = append(out, suffix...)
+	return out, lengthOffset
+}
+
 func TestPatchWifiLocation(t *testing.T) {
 	payload := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
 	body := testFrame(payload)
@@ -69,20 +92,121 @@ func TestPatchWifiLocation(t *testing.T) {
 }

 func TestPatchCellLocation(t *testing.T) {
-	cell := writeLengthDelimited(5, testLocation(300, 400, 25))
-	payload := writeLengthDelimited(22, cell)
-	body := testFrame(payload)
-	c := wlocCoords{Latitude: 22.544577, Longitude: 113.94114, Accuracy: 25}
+	for _, field := range []int{22, 24} {
+		t.Run(fmt.Sprintf("field_%d", field), func(t *testing.T) {
+			cell := writeLengthDelimited(5, testLocation(300, 400, 25))
+			payload := writeLengthDelimited(field, cell)
+			body := testFrame(payload)
+			c := wlocCoords{Latitude: 22.544577, Longitude: 113.94114, Accuracy: 25}
+
+			patched, stats, err := patchWlocBody(body, c)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if stats.Cell != 1 || stats.Locations != 1 {
+				t.Fatalf("unexpected stats: %+v", stats)
+			}
+			if bytes.Equal(patched, body) {
+				t.Fatal("body was not patched")
+			}
+		})
+	}
+}
+
+func TestPatchARPCFramePreservesEnvelopeAndSuffix(t *testing.T) {
+	payload := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
+	suffix := []byte{0xde, 0xad, 0xbe, 0xef}
+	body, lengthOffset := testARPCFrame(payload, suffix)
+	originalPrefix := cloneBytes(body[:lengthOffset])
+	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}

 	patched, stats, err := patchWlocBody(body, c)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if stats.Cell != 1 || stats.Locations != 1 {
+	if stats.WiFi != 1 || stats.Locations != 1 {
 		t.Fatalf("unexpected stats: %+v", stats)
 	}
-	if bytes.Equal(patched, body) {
-		t.Fatal("body was not patched")
+	if !bytes.Equal(patched[:lengthOffset], originalPrefix) {
+		t.Fatal("ARPC metadata changed")
+	}
+	newLength := int(binary.BigEndian.Uint32(patched[lengthOffset : lengthOffset+4]))
+	if newLength == len(payload) {
+		t.Fatal("ARPC payload length was not updated")
+	}
+	if !bytes.Equal(patched[lengthOffset+4+newLength:], suffix) {
+		t.Fatal("ARPC suffix changed")
+	}
+}
+
+func TestPatchARPCPayloadLargerThanUint16(t *testing.T) {
+	padding := writeLengthDelimited(99, bytes.Repeat([]byte{0x7f}, 70_000))
+	location := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
+	payload := append(padding, location...)
+	body, lengthOffset := testARPCFrame(payload, nil)
+	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}
+
+	patched, stats, err := patchWlocBody(body, c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stats.WiFi != 1 || stats.Locations != 1 {
+		t.Fatalf("unexpected stats: %+v", stats)
+	}
+	newLength := int(binary.BigEndian.Uint32(patched[lengthOffset : lengthOffset+4]))
+	if newLength <= 65535 {
+		t.Fatalf("expected 32-bit ARPC payload length, got %d", newLength)
+	}
+	if !bytes.Contains(patched[lengthOffset+4:lengthOffset+4+newLength], padding) {
+		t.Fatal("unknown ARPC payload field changed")
+	}
+}
+
+func TestPatchMarkerFramePreservesPrefixAndSuffix(t *testing.T) {
+	payload := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
+	prefix := []byte{0xaa, 0xbb, 0xcc}
+	suffix := []byte{0xdd, 0xee}
+	var length [2]byte
+	binary.BigEndian.PutUint16(length[:], uint16(len(payload)))
+	body := append(cloneBytes(prefix), wlocMarker...)
+	body = append(body, length[:]...)
+	body = append(body, payload...)
+	body = append(body, suffix...)
+	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}
+
+	patched, stats, err := patchWlocBody(body, c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stats.WiFi != 1 || stats.Locations != 1 {
+		t.Fatalf("unexpected stats: %+v", stats)
+	}
+	lengthOffset := len(prefix) + len(wlocMarker)
+	newLength := int(binary.BigEndian.Uint16(patched[lengthOffset : lengthOffset+2]))
+	if newLength == len(payload) {
+		t.Fatal("marker payload length was not updated")
+	}
+	if !bytes.Equal(patched[:len(prefix)], prefix) {
+		t.Fatal("marker prefix changed")
+	}
+	if !bytes.Equal(patched[lengthOffset+2+newLength:], suffix) {
+		t.Fatal("marker suffix changed")
+	}
+}
+
+func TestPatchBareWlocPayload(t *testing.T) {
+	body := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
+	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}
+
+	patched, stats, err := patchWlocBody(body, c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stats.WiFi != 1 || stats.Locations != 1 {
+		t.Fatalf("unexpected stats: %+v", stats)
+	}
+	if len(patched) > 0 && bytes.HasPrefix(patched, []byte{0, 1, 0, 0}) {
+		t.Fatal("bare payload was unexpectedly wrapped")
 	}
 }