Inject protobuf tag to struct definition in go file.

prototag.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
)

func main() {
	files := os.Args[1:]
	for _, file := range files {
		if err := tag(file); err != nil {
			log.Fatalf("failed to tag %s: %v", file, err)
		}
	}
}

func tag(file string) (err error) {
	cur, err := os.Getwd()
	if err != nil {
		return
	}
	f, err := os.Open(file)
	if err != nil {
		return
	}
	defer f.Close()
	log.Printf("reading file %s", file)

	buf := &bytes.Buffer{}
	var index int
	scanner := bufio.NewScanner(f)
	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		return scanLines(data, atEOF, &index)
	})
	for scanner.Scan() {
		if _, err = buf.Write(scanner.Bytes()); err != nil {
			return
		}
	}

	tmp, err := os.CreateTemp(cur, "prototag-")
	if err != nil {
		return
	}
	defer os.Remove(tmp.Name())
	defer tmp.Close()

	if _, err = buf.WriteTo(tmp); err != nil {
		return
	}

	if err = tmp.Sync(); err != nil {
		return
	}

	return os.Rename(tmp.Name(), file)
}

func scanLines(data []byte, atEOF bool, index *int) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, tagLine(data[0:i+1], index), nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), tagLine(data, index), nil
	}
	// Request more data.
	return 0, nil, nil
}

var commentPrefix = []byte("//")

func tagLine(data []byte, index *int) []byte {
	if bytes.Contains(data, []byte("struct {")) {
		// beginning line of a struct
		*index = 0
	}

	if bytes.Contains(data, []byte("protobuf:")) {
		// skip
		*index++
		return data
	}

	// find out index of the space before `json:`
	i := bytes.Index(data, []byte(" `json:"))
	if i == -1 {
		// line does not contain json tag
		return data
	}

	if bytes.Contains(data[:i], commentPrefix) {
		// json tag in comment
		return data
	}

	if bytes.HasPrefix(data[i:], []byte(" `json:\",inline\"`")) {
		// skip `json:",inline"`
		return data
	}

	*index++
	// cut out the tag name from `json:"xxx:"`
	name := cutName(data[i+1:])
	p := protobufTag(*index, name, data)
	k := bytes.IndexRune(data[i+2:], '`')
	return insertBytes(data, i+2+k, p)
}

var prefixLen = len(` json:"`)

func cutName(data []byte) []byte {
	sub := data[prefixLen:]
	l := 0
	for i, b := range sub {
		if rune(b) == '"' || rune(b) == ',' {
			l = i
			break
		}
	}
	return sub[:l]
}

func insertBytes(data []byte, index int, sub []byte) []byte {
	if index < 0 || index > len(data) {
		return data
	}

	new := make([]byte, len(data)+len(sub))
	copy(new, data[:index])
	copy(new[index:], sub)
	copy(new[index+len(sub):], data[index:])
	return new
}

func protobufTag(index int, name []byte, raw []byte) []byte {
	// remove comment
	if i := bytes.Index(raw, commentPrefix); i != -1 {
		raw = raw[:i]
	}
	buf := bytes.NewBufferString(fmt.Sprintf(` protobuf:"bytes,%d`, index))
	if bytes.Contains(raw, []byte("[]")) {
		_, _ = buf.WriteString(",rep")
	}
	if bytes.Contains(raw, []byte("omitempty")) {
		_, _ = buf.WriteString(",opt")
	}
	_, _ = buf.WriteString(fmt.Sprintf(`,name=%s"`, name))
	return buf.Bytes()
}

$ curl -o prototag.go -sSL0 https://gist.github.com/crazytaxii/15d066bddaa70216c7e228ded379cf8d
$ go build # build prototag
$ ./prototag /path/to/your/file.go # inject protobuf tags