feat: improve wloc response compatibility

Core/wloc_patch.go

@@ -41,6 +41,8 @@ type wireField struct {
 
 var macPattern = regexp.MustCompile(`^[0-9a-fA-F]{1,2}(:[0-9a-fA-F]{1,2}){5}$`)
 
+var wlocMarker = []byte{0, 0, 0, 1, 0, 0}
+
 func minInt(a, b int) int {
 	if a < b {
 		return a
@@ -314,6 +316,99 @@ func patchWlocPayload(payload []byte, c wlocCoords, st *patchStats) ([]byte, boo
 	return out, changed, nil
 }
 
+func parseARPCPayloadBounds(body []byte) (lengthOffset, payloadOffset, payloadEnd int, err error) {
+	if len(body) < 2 {
+		return 0, 0, 0, errors.New("ARPC body too short")
+	}
+
+	offset := 2 // version
+	for range 3 {
+		if offset+2 > len(body) {
+			return 0, 0, 0, errors.New("truncated ARPC string length")
+		}
+		length := int(binary.BigEndian.Uint16(body[offset : offset+2]))
+		offset += 2
+		if length > len(body)-offset {
+			return 0, 0, 0, errors.New("truncated ARPC string")
+		}
+		offset += length
+	}
+
+	const functionAndLengthBytes = 8
+	if offset+functionAndLengthBytes > len(body) {
+		return 0, 0, 0, errors.New("truncated ARPC header")
+	}
+	lengthOffset = offset + 4
+	payloadOffset = lengthOffset + 4
+	payloadLength := uint64(binary.BigEndian.Uint32(body[lengthOffset:payloadOffset]))
+	if payloadLength == 0 || payloadLength > uint64(len(body)-payloadOffset) {
+		return 0, 0, 0, errors.New("invalid ARPC payload length")
+	}
+	return lengthOffset, payloadOffset, payloadOffset + int(payloadLength), nil
+}
+
+func patchARPCFrame(body []byte, c wlocCoords) ([]byte, patchStats, error) {
+	lengthOffset, payloadOffset, payloadEnd, err := parseARPCPayloadBounds(body)
+	if err != nil {
+		return nil, patchStats{}, err
+	}
+
+	var st patchStats
+	payload := body[payloadOffset:payloadEnd]
+	newPayload, changed, err := patchWlocPayload(payload, c, &st)
+	if err != nil {
+		return nil, patchStats{}, err
+	}
+	if !changed || bytes.Equal(newPayload, payload) {
+		return nil, patchStats{}, errors.New("ARPC envelope has no patchable wloc payload")
+	}
+
+	var lenBytes [4]byte
+	binary.BigEndian.PutUint32(lenBytes[:], uint32(len(newPayload)))
+	out := append(cloneBytes(body[:lengthOffset]), lenBytes[:]...)
+	out = append(out, newPayload...)
+	out = append(out, body[payloadEnd:]...)
+	return out, st, nil
+}
+
+func patchMarkerFrame(body []byte, c wlocCoords) ([]byte, patchStats, error) {
+	markerOffset := bytes.Index(body, wlocMarker)
+	if markerOffset < 0 {
+		return nil, patchStats{}, errors.New("wloc marker not found")
+	}
+
+	lengthOffset := markerOffset + len(wlocMarker)
+	payloadOffset := lengthOffset + 2
+	if payloadOffset > len(body) {
+		return nil, patchStats{}, errors.New("truncated marker frame")
+	}
+	payloadLength := int(binary.BigEndian.Uint16(body[lengthOffset:payloadOffset]))
+	if payloadLength == 0 || payloadLength > len(body)-payloadOffset {
+		return nil, patchStats{}, errors.New("invalid marker payload length")
+	}
+	payloadEnd := payloadOffset + payloadLength
+
+	var st patchStats
+	payload := body[payloadOffset:payloadEnd]
+	newPayload, changed, err := patchWlocPayload(payload, c, &st)
+	if err != nil {
+		return nil, patchStats{}, err
+	}
+	if !changed || bytes.Equal(newPayload, payload) {
+		return nil, patchStats{}, errors.New("marker frame has no patchable wloc payload")
+	}
+	if len(newPayload) > 65535 {
+		return nil, patchStats{}, errors.New("patched marker payload too large")
+	}
+
+	var lenBytes [2]byte
+	binary.BigEndian.PutUint16(lenBytes[:], uint16(len(newPayload)))
+	out := append(cloneBytes(body[:lengthOffset]), lenBytes[:]...)
+	out = append(out, newPayload...)
+	out = append(out, body[payloadEnd:]...)
+	return out, st, nil
+}
+
 func patchFrame(body []byte, offset int, c wlocCoords, st *patchStats) ([]byte, patchStats, error) {
 	if len(body) < offset+10 {
 		return nil, *st, fmt.Errorf("body too short: %d, base=%d", len(body), offset)
@@ -352,6 +447,13 @@ func patchFrame(body []byte, offset int, c wlocCoords, st *patchStats) ([]byte,
 }
 
 func patchWlocBody(body []byte, c wlocCoords) ([]byte, patchStats, error) {
+	if out, st, err := patchARPCFrame(body, c); err == nil {
+		return out, st, nil
+	}
+	if out, st, err := patchMarkerFrame(body, c); err == nil {
+		return out, st, nil
+	}
+
 	var st patchStats
 	offsets := []int{0, 2, 4, 6, 8, 10, 12, 14, 16}
 	seen := map[int]bool{}

Core/wloc_patch_test.go

@@ -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")
 	}
 }