scripts/patch_plugin_guest.go

Rama main · réplica de lectura del repositorio original.

//go:build ignore

// Command patch_plugin_guest repairs the guest stubs emitted by
// protoc-gen-go-plugin.
//
// The generator drops the response buffer that the host allocated in guest
// memory: it unmarshals from it and never frees it, so every host call leaks.
// Every committed *_plugin.pb.go therefore frees it after unmarshal. Run this
// from the //go:generate directive right after protoc, otherwise regenerating
// a module silently reintroduces the leak.
//
// Usage: go run scripts/patch_plugin_guest.go -file <file.pb.go> [-file ...]
//
// The targets are passed as flags rather than positionally because `go run`
// would otherwise treat a trailing .go path as another source file to compile.
//
// Idempotent: running it on an already-patched file changes nothing.
package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
)

const (
	header    = "// Code generated by protoc-gen-go-plugin. DO NOT EDIT.\n"
	patchNote = "// Patched: response buffers allocated by the host are freed after unmarshal.\n"

	unpatched = `	if err = response.UnmarshalVT(buf); err != nil {
		return nil, err
	}
`

	patched = `	err = response.UnmarshalVT(buf)
	if ptr != 0 {
		wasm.Free(ptr)
	}
	if err != nil {
		return nil, err
	}
`
)

func main() {
	var files []string

	flag.Func("file", "generated guest stub to patch (repeatable)", func(path string) error {
		files = append(files, path)

		return nil
	})
	flag.Parse()

	if len(files) == 0 {
		fmt.Fprintln(os.Stderr, "usage: patch_plugin_guest.go -file <file.pb.go> [-file ...]")
		os.Exit(2)
	}

	for _, path := range files {
		if err := patch(path); err != nil {
			fmt.Fprintf(os.Stderr, "patch_plugin_guest: %s: %v\n", path, err)
			os.Exit(1)
		}
	}
}

func patch(path string) error {
	content, err := os.ReadFile(path)
	if err != nil {
		return err
	}

	src := string(content)

	if !strings.Contains(src, header) {
		return fmt.Errorf("not a protoc-gen-go-plugin file")
	}

	count := strings.Count(src, unpatched)

	// Nothing to replace is only fine when the file already carries the fix.
	// Otherwise the generator changed its output and this patch no longer
	// matches it — writing the file anyway would advertise a fix that is not
	// there and let the leak back in unnoticed.
	if count == 0 {
		if !strings.Contains(src, patched) {
			return fmt.Errorf("no response buffer unmarshal matched: generator output changed, update the patch")
		}

		fmt.Printf("patch_plugin_guest: %s: already patched\n", path)

		return nil
	}

	if !strings.Contains(src, patchNote) {
		src = strings.Replace(src, header, header+patchNote, 1)
	}

	src = strings.ReplaceAll(src, unpatched, patched)

	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		return err
	}

	fmt.Printf("patch_plugin_guest: %s: freed %d response buffer(s)\n", path, count)

	return nil
}