Fix download reliability with kbsync

cmd/list_versions.go

@@ -63,7 +63,7 @@ func ListVersionsCmd() *cobra.Command {
 					app = lookupResult.App
 				}
 
-				out, err := dependencies.AppStore.ListVersions(appstore.ListVersionsInput{Account: acc, App: app, Platform: platform})
+				out, err := dependencies.AppStore.ListVersions(appstore.ListVersionsInput{Context: cmd.Context(), Account: acc, App: app, Platform: platform})
 				if err != nil {
 					return err
 				}

internal/sap/machine/kbsync.go

@@ -0,0 +1,75 @@
+package machine
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/majd/ipatool/v2/internal/sap/assets"
+)
+
+// GenerateKBSync creates the account and hardware bound FairPlay data required
+// by the bag's ent/download endpoint, without opening a decryption session.
+func GenerateKBSync(ctx context.Context, bundle assets.Bundle, hardwareID []byte, dsid uint64) ([]byte, error) {
+	if ctx == nil {
+		return nil, errors.New("StoreAgent context is nil")
+	}
+
+	if err := ctx.Err(); err != nil {
+		return nil, fmt.Errorf("generate kbsync: %w", err)
+	}
+
+	if dsid == 0 {
+		return nil, errors.New("kbsync requires a nonzero account DSID")
+	}
+
+	image, err := assets.LoadStoreAgent(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("load Apple StoreAgent asset: %w", err)
+	}
+
+	agent, globalContext, err := openStoreAgentGlobal(ctx, bundle, image, hardwareID)
+	if err != nil {
+		return nil, err
+	}
+	defer agent.Close()
+
+	return agent.guest.generateKBSync(storeAgentKBSyncEntry, globalContext, dsid)
+}
+
+func (m *Machine) generateKBSync(entry uint64, globalContext uint32, dsid uint64) ([]byte, error) {
+	m.beginCall()
+	defer m.clearScratch()
+
+	pointerField, err := m.scratch(nil, 8)
+	if err != nil {
+		return nil, err
+	}
+
+	// StoreAgent writes a uint32 length. Zero the upper half so consumeOutput
+	// can share the SAP buffer bounds checking and storage disposal.
+	lengthField, err := m.scratch(nil, 8)
+	if err != nil {
+		return nil, err
+	}
+
+	status, err := m.invoke(entry, uint64(globalContext), dsid, 0, 1, pointerField, lengthField)
+	if err != nil {
+		return nil, fmt.Errorf("generate StoreAgent kbsync: %w", err)
+	}
+
+	output, outputErr := m.consumeOutput(pointerField, lengthField)
+	if int32(status) != 0 {
+		return nil, errors.Join(fmt.Errorf("StoreAgent kbsync returned %d", int32(status)), outputErr)
+	}
+
+	if outputErr != nil {
+		return nil, outputErr
+	}
+
+	if len(output) == 0 {
+		return nil, errors.New("StoreAgent kbsync returned an empty buffer")
+	}
+
+	return output, nil
+}

internal/sap/machine/kbsync_test.go

@@ -0,0 +1,91 @@
+package machine
+
+import (
+	"context"
+	"encoding/binary"
+
+	"github.com/majd/ipatool/v2/internal/sap/assets"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("StoreAgent kbsync", func() {
+	It("pins the kbsync entry in the verified StoreAgent image", func() {
+		Expect(storeAgentKBSyncEntry).To(Equal(storeAgentBase + 0x0c93c0))
+	})
+
+	It("rejects a canceled context before loading assets", func() {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := GenerateKBSync(ctx, assets.Bundle{}, []byte{1}, 123)
+		Expect(err).To(MatchError(ContainSubstring(context.Canceled.Error())))
+	})
+
+	It("rejects missing context or account identity before loading assets", func() {
+		var missingContext context.Context
+		_, err := GenerateKBSync(missingContext, assets.Bundle{}, nil, 123)
+		Expect(err).To(MatchError("StoreAgent context is nil"))
+		_, err = GenerateKBSync(context.Background(), assets.Bundle{}, nil, 0)
+		Expect(err).To(MatchError("kbsync requires a nonzero account DSID"))
+	})
+
+	DescribeTable("copies and disposes guest output with the StoreAgent ABI",
+		func(pointer, length uint64, status uint32, wantError string) {
+			guest := newGinkgoServiceMachine(shimOptions{})
+			data := []byte{0, 4, 0, 1, 42}
+			Expect(guest.engine.MemWrite(heapBase, data)).To(Succeed())
+			var disposed []uint64
+
+			dispose, err := guest.services.addFunction("test.kbsync.dispose", func() error {
+				address, err := guest.services.argument(0)
+				Expect(err).NotTo(HaveOccurred())
+				disposed = append(disposed, address)
+
+				return guest.services.setResult(0)
+			})
+			Expect(err).NotTo(HaveOccurred())
+			guest.entry.dispose = dispose
+
+			entry, err := guest.services.addFunction("test.kbsync.generate", func() error {
+				for index, expected := range []uint64{0x1234, 0x123456789abcdef0, 0, 1} {
+					argument, err := guest.services.argument(index)
+					Expect(err).NotTo(HaveOccurred())
+					Expect(argument).To(Equal(expected))
+				}
+
+				pointerField, err := guest.services.argument(4)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(guest.writeUint64(pointerField, pointer)).To(Succeed())
+				lengthField, err := guest.services.argument(5)
+				Expect(err).NotTo(HaveOccurred())
+				// Only the low 32 bits are written by the real entry point.
+				Expect(guest.engine.MemWrite(lengthField, binary.LittleEndian.AppendUint32(nil, uint32(length)))).To(Succeed())
+
+				return guest.services.setResult(uint64(status))
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			output, err := guest.generateKBSync(entry, 0x1234, 0x123456789abcdef0)
+			if wantError == "" {
+				Expect(err).NotTo(HaveOccurred())
+				Expect(output).To(Equal(data))
+			} else {
+				Expect(err).To(MatchError(ContainSubstring(wantError)))
+				Expect(output).To(BeNil())
+			}
+
+			if pointer != 0 {
+				Expect(disposed).To(Equal([]uint64{pointer}))
+			} else {
+				Expect(disposed).To(BeEmpty())
+			}
+			Expect(guest.scratchCursor).To(BeZero())
+			Expect(guest.engine.MemRead(scratchBase, 32)).To(Equal(make([]byte, 32)))
+		},
+		Entry("valid blob", heapBase, uint64(5), uint32(0), ""),
+		Entry("guest error still releases storage", heapBase, uint64(5), ^uint32(41), "kbsync returned -42"),
+		Entry("oversized blob", heapBase, maxOutputSize+1, uint32(0), "maximum is"),
+		Entry("null pointer", uint64(0), uint64(5), uint32(0), "null output pointer"),
+		Entry("empty blob", heapBase, uint64(0), uint32(0), "empty buffer"),
+	)
+})

internal/sap/machine/storeagent.go

@@ -13,6 +13,7 @@ import (
 const (
 	storeAgentBase         = uint64(0x00001000c0000000)
 	storeAgentGlobalInit   = storeAgentBase + 0x0c5fc0
+	storeAgentKBSyncEntry  = storeAgentBase + 0x0c93c0
 	storeAgentSessionInit  = storeAgentBase + 0x0debd0
 	storeAgentDecryptEntry = storeAgentBase + 0x0ee700
 	storeAgentSessionClose = storeAgentBase + 0x1212d0
@@ -55,63 +56,69 @@ func OpenStoreAgent(ctx context.Context, bundle assets.Bundle, hardwareID, dpInf
 }
 
 func openStoreAgent(ctx context.Context, bundle assets.Bundle, image, hardwareID, dpInfo []byte) (*StoreAgent, error) {
-	if ctx == nil {
-		return nil, errors.New("StoreAgent context is nil")
+	if len(dpInfo) == 0 {
+		return nil, errors.New("StoreAgent dpInfo is empty")
 	}
 
-	if err := ctx.Err(); err != nil {
-		return nil, fmt.Errorf("open StoreAgent runtime: %w", err)
+	agent, globalContext, err := openStoreAgentGlobal(ctx, bundle, image, hardwareID)
+	if err != nil {
+		return nil, err
 	}
 
-	if len(dpInfo) == 0 {
-		return nil, errors.New("StoreAgent dpInfo is empty")
+	agent.session, err = agent.initializeSession(globalContext, dpInfo)
+	if err != nil {
+		_ = agent.Close()
+
+		return nil, err
+	}
+
+	return agent, nil
+}
+
+func openStoreAgentGlobal(ctx context.Context, bundle assets.Bundle, image, hardwareID []byte) (*StoreAgent, uint32, error) {
+	if ctx == nil {
+		return nil, 0, errors.New("StoreAgent context is nil")
+	}
+
+	if err := ctx.Err(); err != nil {
+		return nil, 0, fmt.Errorf("open StoreAgent runtime: %w", err)
 	}
 
 	if err := assets.VerifyStoreAgent(image); err != nil {
-		return nil, fmt.Errorf("verify Apple StoreAgent profile: %w", err)
+		return nil, 0, fmt.Errorf("verify Apple StoreAgent profile: %w", err)
 	}
 
 	hardware, err := hardwareBlock(hardwareID)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	guest, _, err := openRuntime(ctx, bundle, runtimeOptions{
+	guest, exports, err := openRuntime(ctx, bundle, runtimeOptions{
 		extraImages: []imageSpec{{name: "storeagent", data: image, base: storeAgentBase}},
 		shims: shimOptions{
 			zeroReturnAliases: storeAgentZeroReturnAliases,
 		},
 	})
 	if err != nil {
-		return nil, fmt.Errorf("start StoreAgent runtime: %w", err)
+		return nil, 0, fmt.Errorf("start StoreAgent runtime: %w", err)
 	}
 
+	guest.entry.dispose = exports["_jEHf8Xzsv8K"]
+
 	agent := &StoreAgent{
 		guest:        guest,
 		decryptEntry: storeAgentDecryptEntry,
 		closeEntry:   storeAgentSessionClose,
 	}
-	complete := false
-
-	defer func() {
-		if !complete {
-			_ = agent.Close()
-		}
-	}()
 
 	globalContext, err := agent.initializeGlobal(hardware)
 	if err != nil {
-		return nil, err
-	}
+		_ = agent.Close()
 
-	agent.session, err = agent.initializeSession(globalContext, dpInfo)
-	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	complete = true
-
-	return agent, nil
+	return agent, globalContext, nil
 }
 
 func (s *StoreAgent) initializeGlobal(hardware []byte) (uint32, error) {

pkg/appstore/appstore.go

@@ -43,6 +43,8 @@ type appstore struct {
 	searchClient        http.Client[searchResult]
 	purchaseClient      http.Client[purchaseResult]
 	downloadClient      http.Client[downloadResult]
+	kbsyncGenerator     kbsyncGenerator
+	kbsyncCache         kbsyncCache
 	platformClient      http.Client[platformVersionLookupResult]
 	storefrontClient    http.Client[[]byte]
 	bagClient           http.Client[bagResult]
@@ -79,6 +81,7 @@ func NewAppStore(args Args) AppStore {
 		searchClient:        http.NewClient[searchResult](clientArgs),
 		purchaseClient:      http.NewClient[purchaseResult](clientArgs),
 		downloadClient:      http.NewClient[downloadResult](clientArgs),
+		kbsyncGenerator:     defaultKBSyncGenerator,
 		platformClient:      http.NewClient[platformVersionLookupResult](clientArgs),
 		storefrontClient:    http.NewClient[[]byte](clientArgs),
 		bagClient:           http.NewClient[bagResult](http.Args{CookieJar: args.CookieJar, Timeout: http.DefaultAuthenticationTimeout}),

pkg/appstore/appstore_bag.go

@@ -77,6 +77,7 @@ type bagResult struct {
 
 type urlBag struct {
 	AuthEndpoint         string `plist:"authenticateAccount,omitempty"`
+	EntDownloadEndpoint  string `plist:"volumeStoreDownloadProduct,omitempty"`
 	RedownloadEndpoint   string `plist:"redownloadProduct,omitempty"`
 	UpdateEndpoint       string `plist:"updateProduct,omitempty"`
 	SAPSetupEndpoint     string `plist:"sign-sap-setup,omitempty"`

pkg/appstore/appstore_download.go

@@ -63,7 +63,7 @@ func (t *appstore) Download(input DownloadInput) (DownloadOutput, error) {
 		}
 	}
 
-	res, resolvedPlatform, err := t.sendDownloadProduct(input.Account, input.App, guid, externalVersionID, input.Platform)
+	res, resolvedPlatform, err := t.sendDownloadProduct(input.Context, input.Account, input.App, guid, externalVersionID, input.Platform)
 	if err != nil {
 		return DownloadOutput{}, err
 	}