dlv based debugging with VS code and Bazel

dlv.py

#!/usr/bin/env python3
# Copyright 2023-2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrapper script to build binaries with Bazel before debugging them."""
import json
import logging
import os
import random
import socket
import socketserver
import subprocess
import sys
import threading

import bazel_utils

from typing import Any, Callable, List, NoReturn, Optional, Sequence

from bazel_utils import bazel, get_execroot, get_gopath, go, log_command

DEBUG = False

log = logging.getLogger(__name__)


def dlv(*args: str) -> NoReturn:
    """Finds or installs dlv and runs it."""
    # find dlv path
    GOPATH = get_gopath()
    dlv = os.path.join(GOPATH, "bin", "dlv")
    if not os.path.exists(dlv):  # not found, install it
        log.warning("dlv not found at %s -> installing it", dlv)
        env = os.environ.copy()
        del env["GOPROXY"]
        go(
            "install",
            "github.com/go-delve/delve/cmd/dlv@latest",
            env=env,
        )
    if DEBUG:
        args = tuple(add_proxy(list(args)))

    cmd = [dlv]
    cmd.extend(args)
    log_command(*cmd)
    if DEBUG:
        subprocess.check_call(cmd)
        sys.exit(0)
    else:
        os.execlp(dlv, *cmd)


def find_rule(args: List[str]) -> Optional[str]:
    """Tries to identify the target argument.

    Modifies the passed args list!"""
    for i, arg in enumerate(args):
        # extra args
        if arg == "--":
            break
        # option
        if arg.startswith("-"):
            continue
        rule = bazel_utils.identify_rule(arg)
        if rule:
            del args[i]
            return rule
    return None


def find_filter(args: List[str]) -> Optional[str]:
    """Tries to identify the test filter.

    >>> find_filter(["asd", "-test.run", "foo", "bar"])
    'foo'
    >>> find_filter(["asd", "-test.run=foo", "bar"])
    'foo'
    >>> find_filter(["asd", "foo", "bar"])
    >>> find_filter(["asd", "-test.run"])
    Traceback (most recent call last):
    ...
    Exception: Missing -test.run argument
    """
    for i, arg in enumerate(args):
        # other args
        if not arg.startswith("-test.run"):
            continue
        if arg.startswith("-test.run="):
            return arg.split("=", 1)[1]
        if i + 1 >= len(args):
            raise Exception("Missing -test.run argument")
        return args[i + 1]
    return None


def debug(rule_type: str, args: Sequence[str]) -> NoReturn:
    """Builds a Bazel target and debugs it with dlv."""

    # find the rule. Only a single rule per package is supported
    args = list(args)
    rule = find_rule(args)
    target = bazel_utils.find_target(rule_type, rule or ":*", find_filter(args))

    bazel_args = ["--strip=never", "-c", "dbg", target]

    # build the test with debug symbols
    bazel("build", *bazel_args)

    # find the binary location
    path = bazel("cquery", "--output=files", *bazel_args).strip()

    # prefix with the workspace path
    path = str(get_execroot() / path)
    cmd = ["exec", path]
    cmd.extend(args)
    dlv(*cmd)


def add_proxy(args: list[str]) -> list[str]:
    host: Optional[str] = None
    port: Optional[str] = None
    new_port = random.randrange(1025, 65536)
    new_args: list[str] = []
    for arg in args:
        if arg.startswith("--listen="):
            addr = arg.split("=", 1)[1]
            host, port = addr.split(":")
            arg = f"--listen={host}:{new_port}"
        new_args.append(arg)

    if host and port:
        listen_addr = (host, int(port))
        log.info("Redirecting %r to %s", listen_addr, new_port)

        class Proxy(socketserver.StreamRequestHandler):

            @staticmethod
            def decode(s: socket.socket, buf: list[int]) -> Any | None:
                decoder = json.JSONDecoder()
                if len(buf) > 0:
                    try:
                        obj, off = decoder.raw_decode(bytes(buf).decode())
                        del buf[:off]
                        return obj
                    except json.JSONDecodeError:
                        pass
                while True:
                    data = s.recv(16 * 1024)
                    if not data:
                        break
                    buf.extend(data)
                    if buf[0] == 0xA:  # remove leading \n
                        buf.pop(0)
                    try:
                        obj, off = decoder.raw_decode(bytes(buf).decode())
                        del buf[:off]
                        return obj
                    except json.JSONDecodeError:
                        pass
                return None

            def handle(self) -> None:
                out = socket.create_connection((host, new_port))
                log.info("New connection")
                rbuf: list[int] = []
                wbuf: list[int] = []
                while True:
                    obj = Proxy.decode(self.connection, rbuf)
                    if not obj:
                        break
                    log.debug("request: %r", obj)
                    out.sendall(json.dumps(obj).encode() + b"\n")
                    obj = Proxy.decode(out, wbuf)
                    if not obj:
                        break
                    self.wfile.write(json.dumps(obj).encode() + b"\n")

        srv = socketserver.TCPServer(listen_addr, Proxy)
        threading.Thread(target=lambda: srv.serve_forever(), daemon=True).start()
    return new_args


def cmd_test(args: Sequence[str]) -> NoReturn:
    """dlv test command.

    Converts the dlv test into bazel build //foo:bar_test and dlv exec.
    """
    debug("go_test", args)


def cmd_debug(args: Sequence[str]) -> NoReturn:
    """dlv test command.

    Converts the dlv debug into bazel build //foo:bar and dlv exec.
    """
    debug("go_binary", args)


def cmd_exec(args: Sequence[str]) -> NoReturn:
    # debugging prebuilt binary. Pass everything to dlv
    dlv("exec", *args)


def cmd_dap(args: Sequence[str]) -> NoReturn:
    """dlv dap command.
    Uses dlv's Debug Adaptor Protocol.

    We can't support this right now as it passes the actual debug command
    (exec/test/debug/...) via TCP, so we would need to implement a proxy server
    to support it.
    """
    log.error("DAP not supported for Bazel. Assuming make style")
    os.execlp("dlv", "dlv", "dap", *args)


def get_handler_or_die(name: str) -> Callable[[Sequence[str]], NoReturn]:
    HANDLERS = {
        "dap": cmd_dap,
        "debug": cmd_debug,
        "exec": cmd_exec,
        "test": cmd_test,
    }
    handler = HANDLERS.get(name, None)
    if not handler:
        raise RuntimeError(f"Command {name} not supported")
    return handler


def main(argv: Optional[Sequence[str]] = None) -> Optional[int]:
    if not argv:
        argv = sys.argv

    if len(argv) < 2:
        print(argv[0], "<cmd>", "[args...]")
        return 1

    logging.basicConfig(level=logging.DEBUG)
    log_command(*argv)
    cmd = argv[1]

    handler = get_handler_or_die(cmd)

    # no need to pass the command
    return handler(sys.argv[2:])


if __name__ == "__main__":
    sys.exit(main(sys.argv))

settings.json

{
    "go.alternateTools": {
        // go shim to use bazels version
        "go": "${workspaceFolder}/scripts/go.py",
        // dlv shim to build debug binaries with bazel
        "dlv": "${workspaceFolder}/scripts/dlv.py"
    },
    "go.delveConfig": {
        // don't deal with DAP for now
        "debugAdapter": "legacy",
        "substitutePath": [
            // vscode passes absolute paths to dlv, but bazel built binaries only have relative paths
            {
                "from": "${workspaceFolder}/",
                "to": ""
            }
        ]
    },
}