Compare commits

..
Author SHA1 Message Date
moloch-- 08f3997756 Fix sliver.h -> main.h 2026-01-02 21:40:45 -06:00
moloch-- 3aebebecbb Implement RenameImport() 2026-01-02 21:22:19 -06:00
moloch-- 1d7eca8e39 Validate module path 2026-01-02 21:06:14 -06:00
moloch-- 773ee5f867 Rename sliver.go -> main.go 2026-01-02 20:55:26 -06:00
moloch-- 7608a4b0a0 Add GoPackage to ImplantConfig pb 2026-01-02 20:41:35 -06:00
moloch-- 9ecd36b7b6 Implement module rename 2026-01-02 20:33:27 -06:00
moloch-- d8fc5a6df4 Tweak project dir structure 2026-01-02 19:02:07 -06:00
moloch-- 649a6582e0 Add -mod=vendor to GoBuild 2026-01-02 18:53:31 -06:00
moloch-- 00e8a0cfca Move main implant code out of main pkg 2026-01-02 18:50:26 -06:00
moloch-- e9890bc8f9 Set limit to unauthed body req size 2026-01-02 17:32:43 -06:00
Joe 8b798f8307 Merge pull request #2083 from BishopFox/feature/mcp
Experiamental MCP Support
2026-01-02 16:25:24 -07:00
moloch-- d4728af1ea Add mcp logging 2026-01-02 16:26:01 -06:00
moloch-- d1b7d91b1a Add basic file operations to mcp 2026-01-02 12:48:35 -06:00
moloch-- cac6dced34 Tweak mcp 2026-01-02 11:14:32 -06:00
moloch-- a35bf0c1fc Implement mcp stdio option 2026-01-02 07:01:55 -06:00
moloch-- 67616e50c1 Wire rpc to mcp 2026-01-01 22:43:52 -06:00
moloch-- 54de8a2a97 Outline mcp 2026-01-01 22:32:25 -06:00
moloch-- 20916c655c Improve kick-operator command 2026-01-01 11:17:07 -06:00
Joe 6c18201ed6 Merge pull request #2081 from BishopFox/fix/rpc-errors
Improve gRPC Error Codes
2026-01-01 09:57:00 -07:00
moloch-- 29165aff0d Add /artifacts to gitignore 2026-01-01 10:04:53 -06:00
moloch-- 704285f28d Fix release workflow 2026-01-01 09:50:59 -06:00
moloch-- a007872fbc Include version history in release action 2026-01-01 09:48:56 -06:00
190 changed files with 41631 additions and 11464 deletions

No files matched your search

+2
View File
@@ -38,6 +38,8 @@ jobs:
steps:
- name: Check Out Code
uses: actions/checkout@v6
with:
fetch-depth: 0
- name: Go 1.25
uses: actions/setup-go@v6
+4
View File
@@ -2,6 +2,10 @@
/.devcontainer
/.idea
/artifacts
/artifacts/**
/Dockerfile.*
/*.sdk.tar.xz
+12
View File
@@ -73,6 +73,18 @@ func GetConsoleLogsDir() string {
return consoleLogsDir
}
// GetMCPLogsDir - Get the Sliver client MCP logs dir ~/.sliver-client/logs/mcp/
func GetMCPLogsDir() string {
mcpLogsDir := filepath.Join(GetClientLogsDir(), "mcp")
if _, err := os.Stat(mcpLogsDir); os.IsNotExist(err) {
err = os.MkdirAll(mcpLogsDir, 0700)
if err != nil {
log.Fatal(err)
}
}
return mcpLogsDir
}
func assetVersion() string {
appDir := GetRootAppDir()
data, err := os.ReadFile(filepath.Join(appDir, versionFileName))
+3
View File
@@ -69,6 +69,9 @@ func init() {
// other command trees below, such as the implant one.
rootCmd.AddCommand(consoleCmd(con))
// MCP stdio server.
rootCmd.AddCommand(mcpCmd(con))
// Implant.
// The implant command allows users to run commands on slivers from their
// system shell. It makes use of pre-runners for connecting to the server
+94
View File
@@ -0,0 +1,94 @@
package cli
/*
Sliver Implant Framework
Copyright (C) 2023 Bishop Fox
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
import (
"fmt"
"os"
"path/filepath"
"github.com/bishopfox/sliver/client/assets"
"github.com/bishopfox/sliver/client/console"
clientmcp "github.com/bishopfox/sliver/client/mcp"
"github.com/bishopfox/sliver/client/transport"
"github.com/spf13/cobra"
)
func mcpCmd(con *console.SliverClient) *cobra.Command {
mcpCmd := &cobra.Command{
Use: "mcp",
Short: "Start the MCP stdio server",
RunE: func(cmd *cobra.Command, args []string) error {
return runMCP(cmd, con)
},
}
mcpCmd.Flags().String("config", "", "path to client config file")
return mcpCmd
}
func runMCP(cmd *cobra.Command, con *console.SliverClient) error {
configPath, _ := cmd.Flags().GetString("config")
config, err := loadMCPClientConfig(configPath)
if err != nil {
fmt.Printf("%s\n", err)
return nil
}
rpc, ln, err := transport.MTLSConnect(config)
if err != nil {
fmt.Printf("Connection to server failed %s\n", err)
return nil
}
defer ln.Close()
con.Rpc = rpc
go handleConnectionLost(ln)
cfg := clientmcp.DefaultConfig()
if err := clientmcp.ServeStdio(cfg, con.Rpc); err != nil {
fmt.Printf("MCP server error: %s\n", err)
return nil
}
return nil
}
func loadMCPClientConfig(path string) (*assets.ClientConfig, error) {
if path != "" {
configPath := path
if !filepath.IsAbs(configPath) {
if _, err := os.Stat(configPath); err != nil {
configPath = filepath.Join(assets.GetConfigDir(), configPath)
}
}
return assets.ReadConfig(configPath)
}
configs := assets.GetConfigs()
switch len(configs) {
case 0:
return nil, fmt.Errorf("no config files found at %s (use --config)", assets.GetConfigDir())
case 1:
for _, config := range configs {
return config, nil
}
default:
return nil, fmt.Errorf("multiple configs found; use --config to select one")
}
return nil, fmt.Errorf("no config files found")
}
+2 -1
View File
@@ -8,7 +8,7 @@ Shared registration glue, scaffolding, and helpers for Sliver client console com
- `command.go` – Provides helpers for binding Cobra commands, flag sets, and visibility filters into the client command tree.
- `server.go` – Defines the server menu composition, wiring each server-side command group onto the root console structure at runtime.
- `sliver.go` – Builds the implant/Sliver menu, attaching per-target command groups and loading alias/extension manifests before exposure.
- `main.go` – Builds the implant/Sliver menu, attaching per-target command groups and loading alias/extension manifests before exposure.
## Sub-packages
@@ -40,6 +40,7 @@ Shared registration glue, scaffolding, and helpers for Sliver client console com
- `licenses/` – Implements the 'licenses' command group for the Sliver client console.
- `loot/` – Implements the 'loot' command group for the Sliver client console.
- `monitor/` – Implements the 'monitor' command group for the Sliver client console.
- `mcp/` – Implements the 'mcp' command group for the Sliver client console.
- `network/` – Implements the 'network' command group for the Sliver client console.
- `operators/` – Implements the 'operators' command group for the Sliver client console.
- `pivots/` – Implements the 'pivots' command group for the Sliver client console.
+76
View File
@@ -0,0 +1,76 @@
package mcp
import (
"github.com/bishopfox/sliver/client/command/flags"
"github.com/bishopfox/sliver/client/command/help"
"github.com/bishopfox/sliver/client/console"
consts "github.com/bishopfox/sliver/client/constants"
clientmcp "github.com/bishopfox/sliver/client/mcp"
"github.com/rsteube/carapace"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
// Commands returns the `mcp` command and its subcommands.
func Commands(con *console.SliverClient) []*cobra.Command {
defaultConfig := clientmcp.DefaultConfig()
mcpCmd := &cobra.Command{
Use: consts.MCPStr,
Short: "Manage the MCP server",
Long: help.GetHelpFor([]string{consts.MCPStr}),
GroupID: consts.NetworkHelpGroup,
Run: func(cmd *cobra.Command, args []string) {
McpCmd(cmd, con, args)
},
}
startCmd := &cobra.Command{
Use: consts.StartStr,
Short: "Start the MCP server",
Long: help.GetHelpFor([]string{consts.MCPStr, consts.StartStr}),
Run: func(cmd *cobra.Command, args []string) {
McpStartCmd(cmd, con, args)
},
}
flags.Bind("", false, startCmd, func(f *pflag.FlagSet) {
f.String("transport", string(defaultConfig.Transport), "mcp transport (http or sse)")
f.String("listen", defaultConfig.ListenAddress, "listen address for the MCP server")
f.String("name", defaultConfig.ServerName, "server name for MCP initialize")
f.String("version", defaultConfig.ServerVersion, "server version for MCP initialize")
})
flags.BindFlagCompletions(startCmd, func(comp *carapace.ActionMap) {
(*comp)["transport"] = carapace.ActionValues("http", "sse")
})
mcpCmd.AddCommand(startCmd)
stopCmd := &cobra.Command{
Use: consts.StopStr,
Short: "Stop the MCP server",
Long: help.GetHelpFor([]string{consts.MCPStr, consts.StopStr}),
Run: func(cmd *cobra.Command, args []string) {
McpStopCmd(cmd, con, args)
},
}
mcpCmd.AddCommand(stopCmd)
consoleCmd := &cobra.Command{
Use: "console",
Aliases: []string{"repl"},
Short: "Open an MCP console session",
Long: help.GetHelpFor([]string{consts.MCPStr, "console"}),
Run: func(cmd *cobra.Command, args []string) {
McpConsoleCmd(cmd, con, args)
},
}
flags.Bind("", false, consoleCmd, func(f *pflag.FlagSet) {
f.String("transport", string(defaultConfig.Transport), "mcp transport (http or sse)")
f.String("url", "", "mcp server URL (defaults to local server settings)")
})
flags.BindFlagCompletions(consoleCmd, func(comp *carapace.ActionMap) {
(*comp)["transport"] = carapace.ActionValues("http", "sse")
})
mcpCmd.AddCommand(consoleCmd)
return []*cobra.Command{mcpCmd}
}
+132
View File
@@ -0,0 +1,132 @@
package mcp
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"strings"
"github.com/bishopfox/sliver/client/console"
clientmcp "github.com/bishopfox/sliver/client/mcp"
"github.com/bishopfox/sliver/client/version"
mcpclient "github.com/mark3labs/mcp-go/client"
mcpapi "github.com/mark3labs/mcp-go/mcp"
"github.com/spf13/cobra"
)
// McpConsoleCmd connects to an MCP server and starts a minimal console.
func McpConsoleCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
rawTransport, _ := cmd.Flags().GetString("transport")
transport, err := clientmcp.ParseTransport(rawTransport)
if err != nil {
con.PrintErrorf("%s\n", err)
return
}
targetURL, _ := cmd.Flags().GetString("url")
if targetURL == "" {
cfg := clientmcp.GetStatus().Config.WithDefaults()
cfg.Transport = transport
targetURL, err = cfg.EndpointURL()
if err != nil {
con.PrintErrorf("%s\n", err)
return
}
}
client, err := newConsoleClient(transport, targetURL)
if err != nil {
con.PrintErrorf("%s\n", err)
return
}
defer client.Close()
ctx := context.Background()
if err := client.Start(ctx); err != nil {
con.PrintErrorf("%s\n", err)
return
}
_, err = client.Initialize(ctx, mcpapi.InitializeRequest{
Params: mcpapi.InitializeParams{
ProtocolVersion: mcpapi.LATEST_PROTOCOL_VERSION,
ClientInfo: mcpapi.Implementation{
Name: "sliver-mcp-console",
Version: version.Version,
},
Capabilities: mcpapi.ClientCapabilities{},
},
})
if err != nil {
con.PrintErrorf("%s\n", err)
return
}
con.PrintInfof("Connected to MCP server (%s)\n", targetURL)
runConsoleREPL(ctx, client, cmd.InOrStdin(), cmd.OutOrStdout())
}
func newConsoleClient(transport clientmcp.Transport, targetURL string) (*mcpclient.Client, error) {
switch transport {
case clientmcp.TransportHTTP:
return mcpclient.NewStreamableHttpClient(targetURL)
case clientmcp.TransportSSE:
return mcpclient.NewSSEMCPClient(targetURL)
default:
return nil, fmt.Errorf("unsupported transport %q", transport)
}
}
func runConsoleREPL(ctx context.Context, client *mcpclient.Client, in io.Reader, out io.Writer) {
reader := bufio.NewReader(in)
for {
fmt.Fprint(out, "mcp> ")
line, err := reader.ReadString('\n')
if err != nil {
if errors.Is(err, io.EOF) {
fmt.Fprintln(out)
return
}
fmt.Fprintf(out, "read error: %s\n", err)
return
}
line = strings.TrimSpace(line)
switch line {
case "":
continue
case "exit", "quit":
return
case "help", "?":
printConsoleHelp(out)
case "tools":
printConsoleTools(ctx, client, out)
default:
fmt.Fprintf(out, "unknown command: %s\n", line)
}
}
}
func printConsoleHelp(out io.Writer) {
fmt.Fprintln(out, "commands: tools, help, exit")
}
func printConsoleTools(ctx context.Context, client *mcpclient.Client, out io.Writer) {
tools, err := client.ListTools(ctx, mcpapi.ListToolsRequest{})
if err != nil {
fmt.Fprintf(out, "tools error: %s\n", err)
return
}
if len(tools.Tools) == 0 {
fmt.Fprintln(out, "no tools available")
return
}
for _, tool := range tools.Tools {
if tool.Description != "" {
fmt.Fprintf(out, "- %s: %s\n", tool.Name, tool.Description)
continue
}
fmt.Fprintf(out, "- %s\n", tool.Name)
}
}
+37
View File
@@ -0,0 +1,37 @@
package mcp
import (
"time"
"github.com/bishopfox/sliver/client/console"
clientmcp "github.com/bishopfox/sliver/client/mcp"
"github.com/spf13/cobra"
)
// McpCmd prints the current MCP server state.
func McpCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
status := clientmcp.GetStatus()
state := "stopped"
if status.Running {
state = "running"
}
con.Printf("Status: %s\n", state)
con.Printf("Transport: %s\n", status.Config.Transport)
con.Printf("Listen: %s\n", status.Config.ListenAddress)
endpoint, err := status.Config.EndpointURL()
if err == nil {
con.Printf("Endpoint: %s\n", endpoint)
} else {
con.PrintErrorf("Endpoint: %s\n", err)
}
if status.Running && !status.StartedAt.IsZero() {
uptime := time.Since(status.StartedAt).Truncate(time.Second)
con.Printf("Uptime: %s\n", uptime)
}
if status.LastError != "" {
con.PrintErrorf("Last error: %s\n", status.LastError)
}
}
+39
View File
@@ -0,0 +1,39 @@
package mcp
import (
"github.com/bishopfox/sliver/client/console"
slivermcp "github.com/bishopfox/sliver/client/mcp"
"github.com/spf13/cobra"
)
// McpStartCmd starts the local MCP server.
func McpStartCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
rawTransport, _ := cmd.Flags().GetString("transport")
transport, err := slivermcp.ParseTransport(rawTransport)
if err != nil {
con.PrintErrorf("%s\n", err)
return
}
listen, _ := cmd.Flags().GetString("listen")
name, _ := cmd.Flags().GetString("name")
version, _ := cmd.Flags().GetString("version")
cfg := slivermcp.Config{
Transport: transport,
ListenAddress: listen,
ServerName: name,
ServerVersion: version,
}.WithDefaults()
if err := slivermcp.Start(cfg, con.Rpc); err != nil {
con.PrintErrorf("%s\n", err)
return
}
con.PrintInfof("Starting MCP server (%s) on %s\n", cfg.Transport, cfg.ListenAddress)
endpoint, err := cfg.EndpointURL()
if err == nil {
con.PrintInfof("Endpoint: %s\n", endpoint)
}
}
+24
View File
@@ -0,0 +1,24 @@
package mcp
import (
"context"
"time"
"github.com/bishopfox/sliver/client/console"
clientmcp "github.com/bishopfox/sliver/client/mcp"
"github.com/spf13/cobra"
)
const mcpStopTimeout = 5 * time.Second
// McpStopCmd stops the local MCP server.
func McpStopCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
ctx, cancel := context.WithTimeout(context.Background(), mcpStopTimeout)
defer cancel()
if err := clientmcp.Stop(ctx); err != nil {
con.PrintErrorf("%s\n", err)
return
}
con.PrintInfof("MCP server stopped\n")
}
+2
View File
@@ -39,6 +39,7 @@ import (
"github.com/bishopfox/sliver/client/command/jobs"
"github.com/bishopfox/sliver/client/command/licenses"
"github.com/bishopfox/sliver/client/command/loot"
"github.com/bishopfox/sliver/client/command/mcp"
"github.com/bishopfox/sliver/client/command/monitor"
"github.com/bishopfox/sliver/client/command/operators"
"github.com/bishopfox/sliver/client/command/reaction"
@@ -106,6 +107,7 @@ func ServerCommands(con *client.SliverClient, serverCmds func() []*cobra.Command
// C2 Network
bind(consts.NetworkHelpGroup,
jobs.Commands,
mcp.Commands,
websites.Commands,
wireguard.Commands,
c2profiles.Commands,
+1
View File
@@ -276,6 +276,7 @@ const (
WgPortFwdStr = "wg-portfwd"
MonitorStr = "monitor"
MonitorConfigStr = "config"
MCPStr = "mcp"
SSHStr = "ssh"
DLLHijackStr = "dllhijack"
InteractiveStr = "interactive"
+118
View File
@@ -0,0 +1,118 @@
package mcp
import (
"fmt"
"net"
"net/url"
"strings"
"github.com/bishopfox/sliver/client/version"
)
// Transport identifies the MCP server transport to expose.
type Transport string
const (
TransportHTTP Transport = "http"
TransportSSE Transport = "sse"
)
// Config controls the local MCP server settings.
type Config struct {
Transport Transport
ListenAddress string
ServerName string
ServerVersion string
}
// DefaultConfig returns the baseline MCP configuration.
func DefaultConfig() Config {
return Config{
Transport: TransportSSE,
ListenAddress: "127.0.0.1:8080",
ServerName: "Sliver MCP",
ServerVersion: version.Version,
}
}
// WithDefaults fills any empty fields with defaults.
func (c Config) WithDefaults() Config {
defaults := DefaultConfig()
if c.Transport == "" {
c.Transport = defaults.Transport
}
if c.ListenAddress == "" {
c.ListenAddress = defaults.ListenAddress
}
if c.ServerName == "" {
c.ServerName = defaults.ServerName
}
if c.ServerVersion == "" {
c.ServerVersion = defaults.ServerVersion
}
return c
}
// ParseTransport validates an input transport string.
func ParseTransport(raw string) (Transport, error) {
value := strings.ToLower(strings.TrimSpace(raw))
if value == "" {
return DefaultConfig().Transport, nil
}
switch Transport(value) {
case TransportHTTP, TransportSSE:
return Transport(value), nil
default:
return "", fmt.Errorf("unsupported transport %q", raw)
}
}
// Validate ensures the configuration can be used for a server start.
func (c Config) Validate() error {
if c.Transport != TransportHTTP && c.Transport != TransportSSE {
return fmt.Errorf("unsupported transport %q", c.Transport)
}
if c.ListenAddress == "" {
return fmt.Errorf("listen address is required")
}
if _, _, err := net.SplitHostPort(c.ListenAddress); err != nil {
return fmt.Errorf("invalid listen address %q: %w", c.ListenAddress, err)
}
if c.ServerName == "" {
return fmt.Errorf("server name is required")
}
if c.ServerVersion == "" {
return fmt.Errorf("server version is required")
}
return nil
}
// EndpointURL returns the base URL for clients to connect.
func (c Config) EndpointURL() (string, error) {
cfg := c.WithDefaults()
if err := cfg.Validate(); err != nil {
return "", err
}
host, port, err := net.SplitHostPort(cfg.ListenAddress)
if err != nil {
return "", err
}
if host == "" {
host = "127.0.0.1"
}
endpoint := url.URL{
Scheme: "http",
Host: net.JoinHostPort(host, port),
}
switch cfg.Transport {
case TransportHTTP:
endpoint.Path = "/mcp"
case TransportSSE:
endpoint.Path = "/sse"
default:
return "", fmt.Errorf("unsupported transport %q", cfg.Transport)
}
return endpoint.String(), nil
}
File diff suppressed because it is too large. Load diff
+67
View File
@@ -0,0 +1,67 @@
package mcp
import (
"fmt"
"log"
"os"
"path/filepath"
"strings"
"time"
"github.com/bishopfox/sliver/client/assets"
)
const mcpLogTimeFormat = "2006-01-02_15-04-05"
func newMCPLogger() (*log.Logger, *os.File, error) {
logsDir := assets.GetMCPLogsDir()
logPath := filepath.Join(logsDir, fmt.Sprintf("%s.log", time.Now().Format(mcpLogTimeFormat)))
logFile, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o600)
if err != nil {
return nil, nil, err
}
logger := log.New(logFile, "", log.LstdFlags|log.Lshortfile)
return logger, logFile, nil
}
func (m *mcpManager) logf(format string, args ...any) {
if m == nil {
return
}
m.mu.Lock()
logger := m.logger
m.mu.Unlock()
if logger != nil {
logger.Printf(format, args...)
}
}
func (s *SliverMCPServer) logf(format string, args ...any) {
if s == nil || s.logger == nil {
return
}
s.logger.Printf(format, args...)
}
func (s *SliverMCPServer) logToolCall(tool, sessionID, beaconID string, extras ...string) {
if s == nil || s.logger == nil {
return
}
parts := make([]string, 0, 2+len(extras))
if sessionID != "" {
parts = append(parts, fmt.Sprintf("session_id=%s", sessionID))
}
if beaconID != "" {
parts = append(parts, fmt.Sprintf("beacon_id=%s", beaconID))
}
for _, extra := range extras {
if extra != "" {
parts = append(parts, extra)
}
}
if len(parts) == 0 {
s.logger.Printf("mcp tool=%s", tool)
return
}
s.logger.Printf("mcp tool=%s %s", tool, strings.Join(parts, " "))
}
+200
View File
@@ -0,0 +1,200 @@
package mcp
import (
"context"
"errors"
"log"
"net/http"
"os"
"sync"
"time"
"github.com/bishopfox/sliver/protobuf/rpcpb"
mcpserver "github.com/mark3labs/mcp-go/server"
)
var (
ErrAlreadyRunning = errors.New("mcp server already running")
ErrNotRunning = errors.New("mcp server not running")
)
// Status describes the running state of the MCP server.
type Status struct {
Running bool
StartedAt time.Time
Config Config
LastError string
}
type serverTransport interface {
Start(addr string) error
Shutdown(ctx context.Context) error
}
type mcpManager struct {
mu sync.Mutex
cfg Config
running bool
startedAt time.Time
lastErr string
server *SliverMCPServer
transport serverTransport
done chan struct{}
logger *log.Logger
logFile *os.File
}
func newManager() *mcpManager {
return &mcpManager{
cfg: DefaultConfig(),
}
}
var defaultManager = newManager()
// GetStatus returns the current MCP server state.
func GetStatus() Status {
return defaultManager.status()
}
// Start launches the MCP server using the provided configuration.
func Start(cfg Config, rpc rpcpb.SliverRPCClient) error {
return defaultManager.start(cfg, rpc)
}
// ServeStdio runs the MCP server over stdio using the provided configuration.
func ServeStdio(cfg Config, rpc rpcpb.SliverRPCClient) error {
cfg = cfg.WithDefaults()
if err := cfg.Validate(); err != nil {
return err
}
logger, logFile, err := newMCPLogger()
if err != nil {
return err
}
defer logFile.Close()
logger.Printf("starting mcp stdio server")
srv := newServer(cfg, rpc, logger)
err = mcpserver.ServeStdio(srv.server)
if err != nil {
logger.Printf("mcp stdio server error: %v", err)
} else {
logger.Printf("mcp stdio server stopped")
}
return err
}
// Stop shuts down the MCP server.
func Stop(ctx context.Context) error {
return defaultManager.stop(ctx)
}
func (m *mcpManager) status() Status {
m.mu.Lock()
defer m.mu.Unlock()
return Status{
Running: m.running,
StartedAt: m.startedAt,
Config: m.cfg,
LastError: m.lastErr,
}
}
func (m *mcpManager) start(cfg Config, rpc rpcpb.SliverRPCClient) error {
cfg = cfg.WithDefaults()
if err := cfg.Validate(); err != nil {
return err
}
m.mu.Lock()
if m.running {
m.mu.Unlock()
return ErrAlreadyRunning
}
logger, logFile, err := newMCPLogger()
if err != nil {
m.mu.Unlock()
return err
}
mcpServer := newServer(cfg, rpc, logger)
var transport serverTransport
switch cfg.Transport {
case TransportHTTP:
transport = mcpserver.NewStreamableHTTPServer(mcpServer.server)
case TransportSSE:
transport = mcpserver.NewSSEServer(mcpServer.server)
default:
m.mu.Unlock()
logFile.Close()
return errors.New("unsupported transport")
}
m.cfg = cfg
m.running = true
m.startedAt = time.Now()
m.lastErr = ""
m.server = mcpServer
m.transport = transport
m.done = make(chan struct{})
m.logger = logger
m.logFile = logFile
addr := cfg.ListenAddress
done := m.done
m.mu.Unlock()
logger.Printf("starting mcp server transport=%s listen=%s", cfg.Transport, cfg.ListenAddress)
go m.run(addr, transport, done)
return nil
}
func (m *mcpManager) run(addr string, transport serverTransport, done chan struct{}) {
err := transport.Start(addr)
if err != nil && !errors.Is(err, http.ErrServerClosed) {
m.logf("mcp server error: %v", err)
m.mu.Lock()
m.lastErr = err.Error()
m.mu.Unlock()
}
m.logf("mcp server stopped")
m.mu.Lock()
m.running = false
m.transport = nil
m.server = nil
logFile := m.logFile
m.logFile = nil
m.logger = nil
m.mu.Unlock()
if logFile != nil {
logFile.Close()
}
close(done)
}
func (m *mcpManager) stop(ctx context.Context) error {
m.mu.Lock()
if !m.running {
m.mu.Unlock()
return ErrNotRunning
}
transport := m.transport
done := m.done
m.mu.Unlock()
if transport == nil {
return ErrNotRunning
}
m.logf("mcp server shutdown requested")
if err := transport.Shutdown(ctx); err != nil {
return err
}
if done == nil {
return nil
}
select {
case <-done:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
+247
View File
@@ -0,0 +1,247 @@
package mcp
import (
"context"
"log"
"github.com/bishopfox/sliver/protobuf/commonpb"
"github.com/bishopfox/sliver/protobuf/rpcpb"
mcpapi "github.com/mark3labs/mcp-go/mcp"
mcpserver "github.com/mark3labs/mcp-go/server"
)
const (
listSessionsAndBeaconsToolName = "list_sessions_and_beacons"
)
type sessionSummary struct {
ID string `json:"id"`
Name string `json:"name"`
Transport string `json:"transport"`
RemoteAddress string `json:"remote_address"`
Hostname string `json:"hostname"`
Username string `json:"username"`
PID int32 `json:"pid"`
OS string `json:"os"`
Arch string `json:"arch"`
Locale string `json:"locale"`
LastCheckin int64 `json:"last_checkin"`
IsDead bool `json:"is_dead"`
Integrity string `json:"integrity"`
}
type beaconSummary struct {
ID string `json:"id"`
Name string `json:"name"`
Transport string `json:"transport"`
RemoteAddress string `json:"remote_address"`
Hostname string `json:"hostname"`
Username string `json:"username"`
PID int32 `json:"pid"`
OS string `json:"os"`
Arch string `json:"arch"`
Locale string `json:"locale"`
LastCheckin int64 `json:"last_checkin"`
NextCheckin int64 `json:"next_checkin"`
Interval int64 `json:"interval"`
Jitter int64 `json:"jitter"`
TasksCount int64 `json:"tasks_count"`
TasksCountCompleted int64 `json:"tasks_count_completed"`
IsDead bool `json:"is_dead"`
Integrity string `json:"integrity"`
}
type listSessionsAndBeaconsResult struct {
Sessions []sessionSummary `json:"sessions"`
SessionsCount int `json:"sessions_count"`
Beacons []beaconSummary `json:"beacons"`
BeaconsCount int `json:"beacons_count"`
}
// SliverMCPServer wraps the MCP server with Sliver RPC access for handlers.
type SliverMCPServer struct {
Rpc rpcpb.SliverRPCClient
server *mcpserver.MCPServer
logger *log.Logger
}
func newServer(cfg Config, rpc rpcpb.SliverRPCClient, logger *log.Logger) *SliverMCPServer {
base := mcpserver.NewMCPServer(
cfg.ServerName,
cfg.ServerVersion,
mcpserver.WithToolCapabilities(false),
)
listSessionsAndBeaconsTool := mcpapi.NewTool(
listSessionsAndBeaconsToolName,
mcpapi.WithDescription("List active sessions and beacons."),
mcpapi.WithReadOnlyHintAnnotation(true),
mcpapi.WithDestructiveHintAnnotation(false),
mcpapi.WithIdempotentHintAnnotation(true),
)
lsTool := mcpapi.NewTool(
lsToolName,
mcpapi.WithDescription("List the contents of a remote directory."),
mcpapi.WithInputSchema[lsArgs](),
mcpapi.WithReadOnlyHintAnnotation(true),
mcpapi.WithDestructiveHintAnnotation(false),
mcpapi.WithIdempotentHintAnnotation(true),
)
cdTool := mcpapi.NewTool(
cdToolName,
mcpapi.WithDescription("Change the current directory on a remote session or beacon."),
mcpapi.WithInputSchema[cdArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(false),
)
catTool := mcpapi.NewTool(
catToolName,
mcpapi.WithDescription("Download and return the contents of a remote file."),
mcpapi.WithInputSchema[catArgs](),
mcpapi.WithReadOnlyHintAnnotation(true),
mcpapi.WithDestructiveHintAnnotation(false),
mcpapi.WithIdempotentHintAnnotation(true),
)
pwdTool := mcpapi.NewTool(
pwdToolName,
mcpapi.WithDescription("Return the current working directory of a remote session or beacon."),
mcpapi.WithInputSchema[pwdArgs](),
mcpapi.WithReadOnlyHintAnnotation(true),
mcpapi.WithDestructiveHintAnnotation(false),
mcpapi.WithIdempotentHintAnnotation(true),
)
rmTool := mcpapi.NewTool(
rmToolName,
mcpapi.WithDescription("Remove a file or directory on the remote target."),
mcpapi.WithInputSchema[rmArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
mvTool := mcpapi.NewTool(
mvToolName,
mcpapi.WithDescription("Move or rename a file or directory on the remote target."),
mcpapi.WithInputSchema[mvArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
cpTool := mcpapi.NewTool(
cpToolName,
mcpapi.WithDescription("Copy a file on the remote target."),
mcpapi.WithInputSchema[cpArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
mkdirTool := mcpapi.NewTool(
mkdirToolName,
mcpapi.WithDescription("Create a directory on the remote target."),
mcpapi.WithInputSchema[mkdirArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
chmodTool := mcpapi.NewTool(
chmodToolName,
mcpapi.WithDescription("Change file or directory permissions on the remote target."),
mcpapi.WithInputSchema[chmodArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
chownTool := mcpapi.NewTool(
chownToolName,
mcpapi.WithDescription("Change file or directory ownership on the remote target."),
mcpapi.WithInputSchema[chownArgs](),
mcpapi.WithReadOnlyHintAnnotation(false),
mcpapi.WithDestructiveHintAnnotation(true),
)
srv := &SliverMCPServer{
Rpc: rpc,
server: base,
logger: logger,
}
srv.server.AddTool(listSessionsAndBeaconsTool, srv.listSessionsAndBeaconsHandler)
srv.server.AddTool(lsTool, srv.lsHandler)
srv.server.AddTool(cdTool, srv.cdHandler)
srv.server.AddTool(catTool, srv.catHandler)
srv.server.AddTool(pwdTool, srv.pwdHandler)
srv.server.AddTool(rmTool, srv.rmHandler)
srv.server.AddTool(mvTool, srv.mvHandler)
srv.server.AddTool(cpTool, srv.cpHandler)
srv.server.AddTool(mkdirTool, srv.mkdirHandler)
srv.server.AddTool(chmodTool, srv.chmodHandler)
srv.server.AddTool(chownTool, srv.chownHandler)
return srv
}
func (s *SliverMCPServer) listSessionsAndBeaconsHandler(ctx context.Context, _ mcpapi.CallToolRequest) (*mcpapi.CallToolResult, error) {
if s.Rpc == nil {
return mcpapi.NewToolResultError("rpc client not configured"), nil
}
s.logToolCall(listSessionsAndBeaconsToolName, "", "")
sessionsResp, err := s.Rpc.GetSessions(ctx, &commonpb.Empty{})
if err != nil {
return mcpapi.NewToolResultErrorFromErr("failed to list sessions", err), nil
}
beaconsResp, err := s.Rpc.GetBeacons(ctx, &commonpb.Empty{})
if err != nil {
return mcpapi.NewToolResultErrorFromErr("failed to list beacons", err), nil
}
result := listSessionsAndBeaconsResult{
Sessions: make([]sessionSummary, 0, len(sessionsResp.GetSessions())),
Beacons: make([]beaconSummary, 0, len(beaconsResp.GetBeacons())),
}
for _, session := range sessionsResp.GetSessions() {
if session == nil {
continue
}
result.Sessions = append(result.Sessions, sessionSummary{
ID: session.ID,
Name: session.Name,
Transport: session.Transport,
RemoteAddress: session.RemoteAddress,
Hostname: session.Hostname,
Username: session.Username,
PID: session.PID,
OS: session.OS,
Arch: session.Arch,
Locale: session.Locale,
LastCheckin: session.LastCheckin,
IsDead: session.IsDead,
Integrity: session.Integrity,
})
}
for _, beacon := range beaconsResp.GetBeacons() {
if beacon == nil {
continue
}
result.Beacons = append(result.Beacons, beaconSummary{
ID: beacon.ID,
Name: beacon.Name,
Transport: beacon.Transport,
RemoteAddress: beacon.RemoteAddress,
Hostname: beacon.Hostname,
Username: beacon.Username,
PID: beacon.PID,
OS: beacon.OS,
Arch: beacon.Arch,
Locale: beacon.Locale,
LastCheckin: beacon.LastCheckin,
NextCheckin: beacon.NextCheckin,
Interval: beacon.Interval,
Jitter: beacon.Jitter,
TasksCount: beacon.TasksCount,
TasksCountCompleted: beacon.TasksCountCompleted,
IsDead: beacon.IsDead,
Integrity: beacon.Integrity,
})
}
result.SessionsCount = len(result.Sessions)
result.BeaconsCount = len(result.Beacons)
return mcpapi.NewToolResultStructuredOnly(result), nil
}
+8
View File
@@ -154,6 +154,14 @@ else
exit 1
fi
# server / gogo / goname
if go test -tags=server,$TAGS ./server/gogo/goname ; then
:
else
cat ~/.sliver/logs/sliver.log
exit 1
fi
# server / gogo
if go test -tags=server,$TAGS ./server/gogo ; then
:
+11
View File
@@ -68,6 +68,16 @@ require (
tailscale.com v1.90.6
)
require (
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
github.com/invopop/jsonschema v0.13.0 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/spf13/cast v1.7.1 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
)
require (
cloud.google.com/go/compute/metadata v0.7.0 // indirect
filippo.io/edwards25519 v1.1.0 // indirect
@@ -133,6 +143,7 @@ require (
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
github.com/lxn/win v0.0.0-20210218163916-a377121e959e // indirect
github.com/mark3labs/mcp-go v0.43.2
github.com/mattn/go-colorable v0.1.14 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mattn/go-localereader v0.0.1 // indirect
+17
View File
@@ -84,7 +84,11 @@ github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiE
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY=
github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/carapace-sh/carapace v1.9.0 h1:IuEP6YeeK4fhrH7MoOyL0I3GMhvROerkw6YGxkx8Fbw=
github.com/carapace-sh/carapace v1.9.0/go.mod h1:Zs3DpsawpFm+8fEfRVYE9EzMBYH8u7S3Aw9yAAKX0JM=
github.com/carapace-sh/carapace-shlex v1.1.0 h1:58QSJTUxghpa/rEn6+Z9UM/Hy0ipVHDUn9S1mp4i+ZE=
@@ -253,6 +257,8 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2 h1:9K06NfxkBh25x56yVhWWlKFE8YpicaSfHwoV8SFbueA=
github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2/go.mod h1:3A9PQ1cunSDF/1rbTq99Ts4pVnycWg+vlPkfeD2NLFI=
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
@@ -274,6 +280,7 @@ github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
github.com/jsimonetti/rtnetlink v1.4.2 h1:Df9w9TZ3npHTyDn0Ev9e1uzmN2odmXd0QX+J5GTEn90=
github.com/jsimonetti/rtnetlink v1.4.2/go.mod h1:92s6LJdE+1iOrw+F2/RO7LYI2Qd8pPpFNNUYW06gcoM=
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
@@ -306,6 +313,10 @@ github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/lxn/win v0.0.0-20210218163916-a377121e959e h1:H+t6A/QJMbhCSEH5rAuRxh+CtW96g0Or0Fxa9IKr4uc=
github.com/lxn/win v0.0.0-20210218163916-a377121e959e/go.mod h1:KxxjdtRkfNoYDCUP5ryK7XJJNTnpC8atvtmTheChOtk=
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
github.com/mark3labs/mcp-go v0.43.2 h1:21PUSlWWiSbUPQwXIJ5WKlETixpFpq+WBpbMGDSVy/I=
github.com/mark3labs/mcp-go v0.43.2/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw=
github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
@@ -399,6 +410,8 @@ github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6Mwd
github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
github.com/spf13/cobra v1.10.1 h1:lJeBwCfmrnXthfAupyUTzJ/J4Nc1RsHC/mSRU2dll/s=
@@ -455,6 +468,8 @@ github.com/ulikunitz/xz v0.5.15 h1:9DNdB5s+SgV3bQ2ApL10xRc35ck0DuIX/isZvIk+ubY=
github.com/ulikunitz/xz v0.5.15/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
github.com/vishvananda/netns v0.0.5 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zdEY=
github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
github.com/xlab/treeprint v1.2.0 h1:HzHnuAF1plUN2zGlAFHbSQP2qJ0ZAD3XF5XD7OesXRQ=
@@ -463,6 +478,8 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
github.com/yiya1989/sshkrb5 v0.0.1 h1:oAKpXLBeZExjYDV9eTL2KypmXU+5oPrJtqDuR5JbWic=
github.com/yiya1989/sshkrb5 v0.0.1/go.mod h1:r1bHcNa22x8rrg/Kj/cDF5bX6UtZrHiQgRFvj3Mnpl8=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
+1 -1
View File
@@ -6,7 +6,7 @@ Core Go implementation of the Sliver implant runtime and supporting subsystems.
## Go Files
- `sliver.go` – Implant entry point that wires transports, handlers, and the task runner together at startup.
- `main.go` – Implant entry point that wires transports, handlers, and the task runner together at startup.
## Sub-packages
@@ -1,4 +1,4 @@
#include "sliver.h"
#include "main.h"
#ifdef __WIN32
#include <windows.h>
+47
View File
@@ -0,0 +1,47 @@
package main
/*
Sliver Implant Framework
Copyright (C) 2019 Bishop Fox
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
// {{if or .Config.IsSharedLib }}
//#include "main.h"
import "C"
// {{end}}
import (
"github.com/bishopfox/sliver/implant/sliver/runner"
)
// {{if or .Config.IsSharedLib .Config.IsShellcode}}
var isRunning bool = false
// {{range .Config.Exports}}
//export {{.}}
func {{.}} () {
if !isRunning {
isRunning = true
main()
}
}
// {{end}}
// {{end}}
func main() {
runner.Main()
}
File renamed without changes.
@@ -1,4 +1,4 @@
package main
package runner
/*
Sliver Implant Framework
@@ -18,12 +18,6 @@ package main
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
// {{if or .Config.IsSharedLib }}
//#include "sliver.h"
import "C"
// {{end}}
import (
"errors"
@@ -105,26 +99,11 @@ func (serv *sliverService) Execute(args []string, r <-chan svc.ChangeRequest, ch
}
}
}
return
}
// {{end}}
// {{if or .Config.IsSharedLib .Config.IsShellcode}}
var isRunning bool = false
// {{range .Config.Exports}}
//export {{.}}
func {{.}} () {
if !isRunning {
isRunning = true
main()
}
}
// {{end}}
// {{end}}
func main() {
func Main() {
// {{if .Config.Debug}}
log.SetFlags(log.LstdFlags | log.Lshortfile)
File diff suppressed because it is too large. Load diff
+1
View File
@@ -148,6 +148,7 @@ message ImplantConfig {
bool ObfuscateSymbols = 12;
string TemplateName = 13;
bool SGNEnabled = 14;
string GoPackage = 15;
bool IncludeMTLS = 53;
bool IncludeHTTP = 16;
+151 -61
View File
@@ -1,7 +1,7 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.3
// protoc v5.29.3
// protoc-gen-go v1.26.0
// protoc v6.33.2
// source: commonpb/common.proto
package commonpb
@@ -21,16 +21,18 @@ const (
)
type Empty struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
}
func (x *Empty) Reset() {
*x = Empty{}
mi := &file_commonpb_common_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *Empty) String() string {
@@ -41,7 +43,7 @@ func (*Empty) ProtoMessage() {}
func (x *Empty) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[0]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -58,20 +60,23 @@ func (*Empty) Descriptor() ([]byte, []int) {
// Request - Common fields used in all gRPC requests
type Request struct {
state protoimpl.MessageState `protogen:"open.v1"`
Async bool `protobuf:"varint,1,opt,name=Async,proto3" json:"Async,omitempty"`
Timeout int64 `protobuf:"varint,2,opt,name=Timeout,proto3" json:"Timeout,omitempty"`
BeaconID string `protobuf:"bytes,8,opt,name=BeaconID,proto3" json:"BeaconID,omitempty"`
SessionID string `protobuf:"bytes,9,opt,name=SessionID,proto3" json:"SessionID,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Async bool `protobuf:"varint,1,opt,name=Async,proto3" json:"Async,omitempty"`
Timeout int64 `protobuf:"varint,2,opt,name=Timeout,proto3" json:"Timeout,omitempty"`
BeaconID string `protobuf:"bytes,8,opt,name=BeaconID,proto3" json:"BeaconID,omitempty"`
SessionID string `protobuf:"bytes,9,opt,name=SessionID,proto3" json:"SessionID,omitempty"`
}
func (x *Request) Reset() {
*x = Request{}
mi := &file_commonpb_common_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *Request) String() string {
@@ -82,7 +87,7 @@ func (*Request) ProtoMessage() {}
func (x *Request) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[1]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -126,24 +131,26 @@ func (x *Request) GetSessionID() string {
}
// Response - Common fields used in all gRPC responses. Note that the Err field
//
// only used when the implant needs to return an error to the server.
// Client<->Server comms should use normal gRPC error handling.
// only used when the implant needs to return an error to the server.
// Client<->Server comms should use normal gRPC error handling.
type Response struct {
state protoimpl.MessageState `protogen:"open.v1"`
Err string `protobuf:"bytes,1,opt,name=Err,proto3" json:"Err,omitempty"`
Async bool `protobuf:"varint,2,opt,name=Async,proto3" json:"Async,omitempty"`
BeaconID string `protobuf:"bytes,8,opt,name=BeaconID,proto3" json:"BeaconID,omitempty"`
TaskID string `protobuf:"bytes,9,opt,name=TaskID,proto3" json:"TaskID,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Err string `protobuf:"bytes,1,opt,name=Err,proto3" json:"Err,omitempty"`
Async bool `protobuf:"varint,2,opt,name=Async,proto3" json:"Async,omitempty"`
BeaconID string `protobuf:"bytes,8,opt,name=BeaconID,proto3" json:"BeaconID,omitempty"`
TaskID string `protobuf:"bytes,9,opt,name=TaskID,proto3" json:"TaskID,omitempty"`
}
func (x *Response) Reset() {
*x = Response{}
mi := &file_commonpb_common_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *Response) String() string {
@@ -154,7 +161,7 @@ func (*Response) ProtoMessage() {}
func (x *Response) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[2]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -199,18 +206,21 @@ func (x *Response) GetTaskID() string {
// File - A basic file data type
type File struct {
state protoimpl.MessageState `protogen:"open.v1"`
Name string `protobuf:"bytes,1,opt,name=Name,proto3" json:"Name,omitempty"`
Data []byte `protobuf:"bytes,2,opt,name=Data,proto3" json:"Data,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Name string `protobuf:"bytes,1,opt,name=Name,proto3" json:"Name,omitempty"`
Data []byte `protobuf:"bytes,2,opt,name=Data,proto3" json:"Data,omitempty"`
}
func (x *File) Reset() {
*x = File{}
mi := &file_commonpb_common_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *File) String() string {
@@ -221,7 +231,7 @@ func (*File) ProtoMessage() {}
func (x *File) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[3]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -252,23 +262,26 @@ func (x *File) GetData() []byte {
// Process - A basic process data type
type Process struct {
state protoimpl.MessageState `protogen:"open.v1"`
Pid int32 `protobuf:"varint,1,opt,name=Pid,proto3" json:"Pid,omitempty"`
Ppid int32 `protobuf:"varint,2,opt,name=Ppid,proto3" json:"Ppid,omitempty"`
Executable string `protobuf:"bytes,3,opt,name=Executable,proto3" json:"Executable,omitempty"`
Owner string `protobuf:"bytes,4,opt,name=Owner,proto3" json:"Owner,omitempty"`
Architecture string `protobuf:"bytes,7,opt,name=Architecture,proto3" json:"Architecture,omitempty"`
SessionID int32 `protobuf:"varint,5,opt,name=SessionID,proto3" json:"SessionID,omitempty"`
CmdLine []string `protobuf:"bytes,6,rep,name=CmdLine,proto3" json:"CmdLine,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Pid int32 `protobuf:"varint,1,opt,name=Pid,proto3" json:"Pid,omitempty"`
Ppid int32 `protobuf:"varint,2,opt,name=Ppid,proto3" json:"Ppid,omitempty"`
Executable string `protobuf:"bytes,3,opt,name=Executable,proto3" json:"Executable,omitempty"`
Owner string `protobuf:"bytes,4,opt,name=Owner,proto3" json:"Owner,omitempty"`
Architecture string `protobuf:"bytes,7,opt,name=Architecture,proto3" json:"Architecture,omitempty"`
SessionID int32 `protobuf:"varint,5,opt,name=SessionID,proto3" json:"SessionID,omitempty"`
CmdLine []string `protobuf:"bytes,6,rep,name=CmdLine,proto3" json:"CmdLine,omitempty"`
}
func (x *Process) Reset() {
*x = Process{}
mi := &file_commonpb_common_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *Process) String() string {
@@ -279,7 +292,7 @@ func (*Process) ProtoMessage() {}
func (x *Process) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[4]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -345,18 +358,21 @@ func (x *Process) GetCmdLine() []string {
// EnvVar - Environment variable K/V
type EnvVar struct {
state protoimpl.MessageState `protogen:"open.v1"`
Key string `protobuf:"bytes,1,opt,name=Key,proto3" json:"Key,omitempty"`
Value string `protobuf:"bytes,2,opt,name=Value,proto3" json:"Value,omitempty"`
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Key string `protobuf:"bytes,1,opt,name=Key,proto3" json:"Key,omitempty"`
Value string `protobuf:"bytes,2,opt,name=Value,proto3" json:"Value,omitempty"`
}
func (x *EnvVar) Reset() {
*x = EnvVar{}
mi := &file_commonpb_common_proto_msgTypes[5]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_commonpb_common_proto_msgTypes[5]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *EnvVar) String() string {
@@ -367,7 +383,7 @@ func (*EnvVar) ProtoMessage() {}
func (x *EnvVar) ProtoReflect() protoreflect.Message {
mi := &file_commonpb_common_proto_msgTypes[5]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -453,7 +469,7 @@ func file_commonpb_common_proto_rawDescGZIP() []byte {
}
var file_commonpb_common_proto_msgTypes = make([]protoimpl.MessageInfo, 6)
var file_commonpb_common_proto_goTypes = []any{
var file_commonpb_common_proto_goTypes = []interface{}{
(*Empty)(nil), // 0: commonpb.Empty
(*Request)(nil), // 1: commonpb.Request
(*Response)(nil), // 2: commonpb.Response
@@ -474,6 +490,80 @@ func file_commonpb_common_proto_init() {
if File_commonpb_common_proto != nil {
return
}
if !protoimpl.UnsafeEnabled {
file_commonpb_common_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*Empty); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_commonpb_common_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*Request); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_commonpb_common_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*Response); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_commonpb_common_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*File); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_commonpb_common_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*Process); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_commonpb_common_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*EnvVar); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
+37 -18
View File
@@ -1,7 +1,7 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.3
// protoc v5.29.3
// protoc-gen-go v1.26.0
// protoc v6.33.2
// source: dnspb/dns.proto
package dnspb
@@ -87,27 +87,32 @@ func (DNSMessageType) EnumDescriptor() ([]byte, []int) {
return file_dnspb_dns_proto_rawDescGZIP(), []int{0}
}
// NOTE: DNS is very space sensitive so certain fields are re-purposed
// depending on the DNSMessageType as noted below:
//
// [Type TOTP]: ID field is used for the TOTP code
//NOTE: DNS is very space sensitive so certain fields are re-purposed
//depending on the DNSMessageType as noted below:
//
//[Type TOTP]: ID field is used for the TOTP code
//
type DNSMessage struct {
state protoimpl.MessageState `protogen:"open.v1"`
Type DNSMessageType `protobuf:"varint,1,opt,name=Type,proto3,enum=dnspb.DNSMessageType" json:"Type,omitempty"`
ID uint32 `protobuf:"varint,2,opt,name=ID,proto3" json:"ID,omitempty"` // 8 bit message id + 24 bit dns session ID
Start uint32 `protobuf:"varint,3,opt,name=Start,proto3" json:"Start,omitempty"` // Bytes start at
Stop uint32 `protobuf:"varint,4,opt,name=Stop,proto3" json:"Stop,omitempty"` // Bytes stop at
Size uint32 `protobuf:"varint,5,opt,name=Size,proto3" json:"Size,omitempty"` // Total size
Data []byte `protobuf:"bytes,6,opt,name=Data,proto3" json:"Data,omitempty"` // Actual data
unknownFields protoimpl.UnknownFields
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Type DNSMessageType `protobuf:"varint,1,opt,name=Type,proto3,enum=dnspb.DNSMessageType" json:"Type,omitempty"`
ID uint32 `protobuf:"varint,2,opt,name=ID,proto3" json:"ID,omitempty"` // 8 bit message id + 24 bit dns session ID
Start uint32 `protobuf:"varint,3,opt,name=Start,proto3" json:"Start,omitempty"` // Bytes start at
Stop uint32 `protobuf:"varint,4,opt,name=Stop,proto3" json:"Stop,omitempty"` // Bytes stop at
Size uint32 `protobuf:"varint,5,opt,name=Size,proto3" json:"Size,omitempty"` // Total size
Data []byte `protobuf:"bytes,6,opt,name=Data,proto3" json:"Data,omitempty"` // Actual data
}
func (x *DNSMessage) Reset() {
*x = DNSMessage{}
mi := &file_dnspb_dns_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
if protoimpl.UnsafeEnabled {
mi := &file_dnspb_dns_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *DNSMessage) String() string {
@@ -118,7 +123,7 @@ func (*DNSMessage) ProtoMessage() {}
func (x *DNSMessage) ProtoReflect() protoreflect.Message {
mi := &file_dnspb_dns_proto_msgTypes[0]
if x != nil {
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
@@ -218,7 +223,7 @@ func file_dnspb_dns_proto_rawDescGZIP() []byte {
var file_dnspb_dns_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
var file_dnspb_dns_proto_msgTypes = make([]protoimpl.MessageInfo, 1)
var file_dnspb_dns_proto_goTypes = []any{
var file_dnspb_dns_proto_goTypes = []interface{}{
(DNSMessageType)(0), // 0: dnspb.DNSMessageType
(*DNSMessage)(nil), // 1: dnspb.DNSMessage
}
@@ -236,6 +241,20 @@ func file_dnspb_dns_proto_init() {
if File_dnspb_dns_proto != nil {
return
}
if !protoimpl.UnsafeEnabled {
file_dnspb_dns_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
switch v := v.(*DNSMessage); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
+3 -3
View File
@@ -1,7 +1,7 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.3
// protoc v5.29.3
// protoc-gen-go v1.26.0
// protoc v6.33.2
// source: rpcpb/services.proto
package rpcpb
@@ -724,7 +724,7 @@ var file_rpcpb_services_proto_rawDesc = []byte{
0x70, 0x63, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
}
var file_rpcpb_services_proto_goTypes = []any{
var file_rpcpb_services_proto_goTypes = []interface{}{
(*commonpb.Empty)(nil), // 0: commonpb.Empty
(*clientpb.ClientLogData)(nil), // 1: clientpb.ClientLogData
(*sliverpb.KillReq)(nil), // 2: sliverpb.KillReq
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+3 -3
View File
@@ -277,7 +277,7 @@ func SetupGoPath(goPathSrc string, includeDNS bool) error {
}
sliverpbGoSrc = xorPBRawBytes(sliverpbGoSrc)
sliverpbGoSrc = stripSliverpb(sliverpbGoSrc)
sliverpbDir := filepath.Join(goPathSrc, "github.com", "bishopfox", "sliver", "protobuf", "sliverpb")
sliverpbDir := filepath.Join(goPathSrc, "protobuf", "sliverpb")
setupLog.Infof("Creating sliverpb directory: %s", sliverpbDir)
os.MkdirAll(sliverpbDir, 0700)
os.WriteFile(filepath.Join(sliverpbDir, "sliver.pb.go"), sliverpbGoSrc, 0600)
@@ -290,7 +290,7 @@ func SetupGoPath(goPathSrc string, includeDNS bool) error {
return err
}
commonpbSrc = xorPBRawBytes(commonpbSrc)
commonpbDir := filepath.Join(goPathSrc, "github.com", "bishopfox", "sliver", "protobuf", "commonpb")
commonpbDir := filepath.Join(goPathSrc, "protobuf", "commonpb")
os.MkdirAll(commonpbDir, 0700)
os.WriteFile(filepath.Join(commonpbDir, "common.pb.go"), commonpbSrc, 0600)
@@ -302,7 +302,7 @@ func SetupGoPath(goPathSrc string, includeDNS bool) error {
return err
}
dnspbSrc = xorPBRawBytes(dnspbSrc)
dnspbDir := filepath.Join(goPathSrc, "github.com", "bishopfox", "sliver", "protobuf", "dnspb")
dnspbDir := filepath.Join(goPathSrc, "protobuf", "dnspb")
os.MkdirAll(dnspbDir, 0700)
os.WriteFile(filepath.Join(dnspbDir, "dns.pb.go"), dnspbSrc, 0600)
}
+5 -1
View File
@@ -69,6 +69,7 @@ var (
const (
DefaultMaxBodyLength = 2 * 1024 * 1024 * 1024 // 2Gb
DefaultMaxUnauthBodyLength = 8 * 1024 * 1024 // 8Mb
DefaultHTTPTimeout = time.Minute
DefaultLongPollTimeout = time.Second
DefaultLongPollJitter = time.Second
@@ -564,7 +565,10 @@ func (s *SliverHTTPC2) anonymousHandler(resp http.ResponseWriter, req *http.Requ
func (s *SliverHTTPC2) startSessionHandler(resp http.ResponseWriter, req *http.Request, encoder sliverEncoders.Encoder) {
httpLog.Debug("Start http session request")
body, err := io.ReadAll(req.Body)
body, err := io.ReadAll(&io.LimitedReader{
R: req.Body,
N: int64(DefaultMaxUnauthBodyLength),
})
if err != nil {
httpLog.Errorf("Failed to read body %s", err)
s.defaultHandler(resp, req)
+94 -2
View File
@@ -28,12 +28,15 @@ import (
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"github.com/spf13/cobra"
consts "github.com/bishopfox/sliver/client/constants"
"github.com/bishopfox/sliver/protobuf/clientpb"
"github.com/bishopfox/sliver/server/certs"
"github.com/bishopfox/sliver/server/console/forms"
"github.com/bishopfox/sliver/server/core"
"github.com/bishopfox/sliver/server/db"
"github.com/bishopfox/sliver/server/db/models"
@@ -172,14 +175,54 @@ func NewOperatorConfig(operatorName string, lhost string, lport uint16, permissi
return json.Marshal(config)
}
func kickOperatorCmd(cmd *cobra.Command, _ []string) {
func kickOperatorCmd(cmd *cobra.Command, args []string) {
operator, _ := cmd.Flags().GetString("name")
if shouldPromptKickOperator(cmd, args) {
names, err := operatorNames()
if err != nil {
fmt.Printf(Warn+"Failed to list operators: %v\n", err)
return
}
if len(names) == 0 {
fmt.Printf(Warn + "No operators found.\n")
return
}
if err := forms.SelectOperator("Select operator to kick", names, &operator); err != nil {
if errors.Is(err, forms.ErrUserAborted) {
return
}
fmt.Printf(Warn+"Operator selection failed: %v\n", err)
return
}
if err := cmd.Flags().Set("name", operator); err != nil {
fmt.Printf(Warn+"Failed to set operator name: %v\n", err)
return
}
}
operator = strings.TrimSpace(operator)
if operator == "" {
fmt.Printf(Warn + "Operator name required (use --name).\n")
return
}
exists, err := operatorExists(operator)
if err != nil {
fmt.Printf(Warn+"Failed to lookup operator %s: %v\n", operator, err)
return
}
if !exists {
fmt.Printf(Warn+"Operator %s does not exist.\n", operator)
return
}
fmt.Printf(Info+"Removing auth token(s) for %s, please wait ... \n", operator)
err := db.Session().Where(&models.Operator{
err = db.Session().Where(&models.Operator{
Name: operator,
}).Delete(&models.Operator{}).Error
if err != nil {
fmt.Printf(Warn+"Failed to remove operator %s: %v\n", operator, err)
return
}
transport.ClearTokenCache()
@@ -192,6 +235,55 @@ func kickOperatorCmd(cmd *cobra.Command, _ []string) {
fmt.Printf(Info+"Operator %s has been kicked out.\n", operator)
}
func shouldPromptKickOperator(cmd *cobra.Command, args []string) bool {
if len(args) != 0 {
return false
}
return cmd.Flags().NFlag() == 0
}
func operatorNames() ([]string, error) {
operators, err := db.OperatorAll()
if err != nil {
return nil, err
}
names := make([]string, 0, len(operators))
seen := make(map[string]struct{}, len(operators))
for _, operator := range operators {
if operator == nil {
continue
}
name := strings.TrimSpace(operator.Name)
if name == "" {
continue
}
if _, ok := seen[name]; ok {
continue
}
seen[name] = struct{}{}
names = append(names, name)
}
sort.Strings(names)
return names, nil
}
func operatorExists(name string) (bool, error) {
if strings.TrimSpace(name) == "" {
return false, nil
}
err := db.Session().Where(&models.Operator{
Name: name,
}).First(&models.Operator{}).Error
if err == nil {
return true, nil
}
if errors.Is(err, db.ErrRecordNotFound) {
return false, nil
}
return false, err
}
func startMultiplayerModeCmd(cmd *cobra.Command, _ []string) {
lhost, _ := cmd.Flags().GetString("lhost")
lport, _ := cmd.Flags().GetUint16("lport")
+74
View File
@@ -0,0 +1,74 @@
package forms
import (
"errors"
"strings"
"github.com/charmbracelet/huh"
)
const defaultSelectHeight = 10
// ErrUserAborted exposes the huh user-abort sentinel for callers.
var ErrUserAborted = huh.ErrUserAborted
// SelectOperator prompts for a single operator selection.
func SelectOperator(title string, operators []string, value *string) error {
if value == nil {
return errors.New("operator selection value is required")
}
if len(operators) == 0 {
return errors.New("operator selection options are required")
}
ensureSelectedValue(operators, value)
field := huh.NewSelect[string]().
Title(title).
Options(makeStringOptions(operators)...).
Height(listHeight(len(operators))).
Value(value).
Validate(func(val string) error {
if strings.TrimSpace(val) == "" {
return errors.New("selection required")
}
return nil
})
form := huh.NewForm(huh.NewGroup(field))
return form.Run()
}
func ensureSelectedValue(options []string, value *string) {
if value == nil || len(options) == 0 {
return
}
if strings.TrimSpace(*value) == "" {
*value = options[0]
return
}
for _, option := range options {
if option == *value {
return
}
}
*value = options[0]
}
func makeStringOptions(options []string) []huh.Option[string] {
converted := make([]huh.Option[string], 0, len(options))
for _, option := range options {
converted = append(converted, huh.NewOption(option, option))
}
return converted
}
func listHeight(count int) int {
if count < 1 {
return 1
}
if count > defaultSelectHeight {
return defaultSelectHeight
}
return count
}
+23 -3
View File
@@ -40,9 +40,11 @@ import (
"github.com/bishopfox/sliver/server/db/models"
"github.com/bishopfox/sliver/server/encoders"
"github.com/bishopfox/sliver/server/gogo"
"github.com/bishopfox/sliver/server/gogo/goname"
"github.com/bishopfox/sliver/server/log"
"github.com/bishopfox/sliver/util"
utilEncoders "github.com/bishopfox/sliver/util/encoders"
"golang.org/x/mod/module"
"golang.org/x/text/cases"
"golang.org/x/text/language"
)
@@ -360,7 +362,7 @@ func renderSliverGoCode(name string, build *clientpb.ImplantBuild, config *clien
return "", err
}
sliverPkgDir := filepath.Join(srcDir, "github.com", "bishopfox", "sliver") // "main"
sliverPkgDir := filepath.Join(srcDir) // "main"
err = os.MkdirAll(sliverPkgDir, 0700)
if err != nil {
return "", nil
@@ -380,14 +382,14 @@ func renderSliverGoCode(name string, build *clientpb.ImplantBuild, config *clien
sliverGoCode := string(sliverGoCodeRaw)
// Skip dllmain files for anything non windows
if f.Name() == "sliver.c" || f.Name() == "sliver.h" {
if f.Name() == "main.c" || f.Name() == "main.h" {
if !config.IsSharedLib && !config.IsShellcode {
return nil
}
}
var sliverCodePath string
if f.Name() == "sliver.go" || f.Name() == "sliver.c" || f.Name() == "sliver.h" {
if f.Name() == "main.go" || f.Name() == "main.c" || f.Name() == "main.h" {
sliverCodePath = filepath.Join(sliverPkgDir, f.Name())
} else {
sliverCodePath = filepath.Join(sliverPkgDir, "implant", fsPath)
@@ -521,6 +523,24 @@ func renderSliverGoCode(name string, build *clientpb.ImplantBuild, config *clien
}
buildLog.Debugf("Created %s", goModPath)
if !config.Debug {
goPackage := "github.com/google/gvisor"
if config.GoPackage != "" {
err := module.CheckPath(config.GoPackage)
if err != nil {
buildLog.Warnf("Invalid Go package path '%s', using default '%s'", config.GoPackage, goPackage)
} else {
goPackage = config.GoPackage
}
}
result, err := goname.RenameModule(sliverPkgDir, goPackage)
if err != nil {
buildLog.Errorf("Failed to rename module: %s", err)
return "", err
}
buildLog.Infof("Renamed module: %s -> %s (go.mod updated: %v)", result.OldModule, result.NewModule, result.GoModUpdated)
}
return sliverPkgDir, nil
}
+1
View File
@@ -183,6 +183,7 @@ func GoBuild(config GoConfig, src string, dest string, buildmode string, tags []
var goCommand = []string{"build"}
goCommand = append(goCommand, "-trimpath") // remove absolute paths from any compiled binary
goCommand = append(goCommand, "-mod=vendor")
if 0 < len(tags) {
goCommand = append(goCommand, "-tags")
+161
View File
@@ -0,0 +1,161 @@
package goname
import (
"bytes"
"errors"
"fmt"
"go/format"
"go/parser"
"go/token"
"io/fs"
"os"
"path/filepath"
"strconv"
"strings"
"golang.org/x/mod/modfile"
)
// Result captures changes made by RenameModule.
type Result struct {
OldModule string
NewModule string
GoModUpdated bool
FilesUpdated int
ImportsUpdated int
}
// RenameModule updates go.mod and rewrites import paths within the module tree.
func RenameModule(dir, newModule string) (*Result, error) {
newModule = strings.TrimSpace(newModule)
if newModule == "" {
return nil, errors.New("new module path is required")
}
goModPath := filepath.Join(dir, "go.mod")
modBytes, err := os.ReadFile(goModPath)
if err != nil {
return nil, fmt.Errorf("read go.mod: %w", err)
}
modFile, err := modfile.Parse(goModPath, modBytes, nil)
if err != nil {
return nil, fmt.Errorf("parse go.mod: %w", err)
}
if modFile.Module == nil || modFile.Module.Mod.Path == "" {
return nil, errors.New("go.mod is missing a module path")
}
oldModule := modFile.Module.Mod.Path
result := &Result{
OldModule: oldModule,
NewModule: newModule,
}
if oldModule == newModule {
return result, nil
}
if err := modFile.AddModuleStmt(newModule); err != nil {
return result, fmt.Errorf("update module path: %w", err)
}
formatted, err := modFile.Format()
if err != nil {
return result, fmt.Errorf("format go.mod: %w", err)
}
if err := writeFilePreservePerm(goModPath, formatted); err != nil {
return result, fmt.Errorf("write go.mod: %w", err)
}
result.GoModUpdated = true
if err := rewriteImports(dir, oldModule, newModule, result); err != nil {
return result, err
}
return result, nil
}
func rewriteImports(dir, oldModule, newModule string, result *Result) error {
return filepath.WalkDir(dir, func(path string, entry fs.DirEntry, err error) error {
if err != nil {
return err
}
if entry.IsDir() {
name := entry.Name()
if name == "vendor" || name == ".git" {
return fs.SkipDir
}
return nil
}
if filepath.Ext(path) != ".go" {
return nil
}
changed, importsUpdated, err := rewriteFileImports(path, oldModule, newModule)
if err != nil {
return err
}
if changed {
result.FilesUpdated++
result.ImportsUpdated += importsUpdated
}
return nil
})
}
func rewriteFileImports(path, oldModule, newModule string) (bool, int, error) {
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
if err != nil {
return false, 0, fmt.Errorf("parse %s: %w", path, err)
}
changed := false
importsUpdated := 0
for _, imp := range file.Imports {
importPath, err := strconv.Unquote(imp.Path.Value)
if err != nil {
return false, 0, fmt.Errorf("parse import in %s: %w", path, err)
}
if !matchesModule(importPath, oldModule) {
continue
}
newPath := newModule + strings.TrimPrefix(importPath, oldModule)
if newPath == importPath {
continue
}
imp.Path.Value = strconv.Quote(newPath)
changed = true
importsUpdated++
}
if !changed {
return false, 0, nil
}
var buf bytes.Buffer
if err := format.Node(&buf, fset, file); err != nil {
return false, 0, fmt.Errorf("format %s: %w", path, err)
}
if err := writeFilePreservePerm(path, buf.Bytes()); err != nil {
return false, 0, fmt.Errorf("write %s: %w", path, err)
}
return true, importsUpdated, nil
}
func matchesModule(importPath, modulePath string) bool {
return importPath == modulePath || strings.HasPrefix(importPath, modulePath+"/")
}
func writeFilePreservePerm(path string, data []byte) error {
perm := fs.FileMode(0644)
if info, err := os.Stat(path); err == nil {
perm = info.Mode().Perm()
}
return os.WriteFile(path, data, perm)
}
+139
View File
@@ -0,0 +1,139 @@
package goname
import (
"go/parser"
"go/token"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"golang.org/x/mod/modfile"
)
func TestRenameModule(t *testing.T) {
fixture := filepath.Join("testdata", "basic")
tmp := t.TempDir()
if err := copyDir(fixture, tmp); err != nil {
t.Fatalf("copy fixture: %v", err)
}
result, err := RenameModule(tmp, "example.net/new")
if err != nil {
t.Fatalf("RenameModule: %v", err)
}
if result.OldModule != "example.com/old" {
t.Fatalf("old module mismatch: %s", result.OldModule)
}
if result.NewModule != "example.net/new" {
t.Fatalf("new module mismatch: %s", result.NewModule)
}
if !result.GoModUpdated {
t.Fatalf("expected go.mod update")
}
if result.FilesUpdated != 2 {
t.Fatalf("files updated mismatch: %d", result.FilesUpdated)
}
if result.ImportsUpdated != 3 {
t.Fatalf("imports updated mismatch: %d", result.ImportsUpdated)
}
modBytes, err := os.ReadFile(filepath.Join(tmp, "go.mod"))
if err != nil {
t.Fatalf("read go.mod: %v", err)
}
modFile, err := modfile.Parse("go.mod", modBytes, nil)
if err != nil {
t.Fatalf("parse go.mod: %v", err)
}
if modFile.Module == nil || modFile.Module.Mod.Path != "example.net/new" {
t.Fatalf("module path not updated: %#v", modFile.Module)
}
mainImports := importPaths(t, filepath.Join(tmp, "main.go"))
for _, want := range []string{"example.net/new", "example.net/new/pkg/foo", "example.com/oldish/pkg"} {
if !mainImports[want] {
t.Fatalf("main.go missing import: %s", want)
}
}
if mainImports["example.com/old"] || mainImports["example.com/old/pkg/foo"] {
t.Fatalf("main.go still references old module")
}
fooImports := importPaths(t, filepath.Join(tmp, "pkg", "foo", "foo.go"))
if !fooImports["example.net/new/pkg/bar"] {
t.Fatalf("foo.go import not updated")
}
vendorImports := importPaths(t, filepath.Join(tmp, "vendor", "example.com", "old", "vendor.go"))
if !vendorImports["example.com/old/pkg/foo"] {
t.Fatalf("vendor import should remain unchanged")
}
mainBytes, err := os.ReadFile(filepath.Join(tmp, "main.go"))
if err != nil {
t.Fatalf("read main.go: %v", err)
}
if !strings.Contains(string(mainBytes), "const example = \"example.com/old/pkg/foo\"") {
t.Fatalf("string literal should remain unchanged")
}
}
func TestRenameModuleMissingGoMod(t *testing.T) {
_, err := RenameModule(t.TempDir(), "example.com/new")
if err == nil {
t.Fatalf("expected error for missing go.mod")
}
}
func importPaths(t *testing.T, path string) map[string]bool {
t.Helper()
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, path, nil, 0)
if err != nil {
t.Fatalf("parse %s: %v", path, err)
}
imports := make(map[string]bool)
for _, imp := range file.Imports {
pathValue, err := strconv.Unquote(imp.Path.Value)
if err != nil {
t.Fatalf("unquote import in %s: %v", path, err)
}
imports[pathValue] = true
}
return imports
}
func copyDir(src, dst string) error {
return filepath.WalkDir(src, func(path string, entry os.DirEntry, err error) error {
if err != nil {
return err
}
rel, err := filepath.Rel(src, path)
if err != nil {
return err
}
target := filepath.Join(dst, rel)
if entry.IsDir() {
info, err := entry.Info()
if err != nil {
return err
}
return os.MkdirAll(target, info.Mode().Perm())
}
data, err := os.ReadFile(path)
if err != nil {
return err
}
info, err := entry.Info()
if err != nil {
return err
}
return os.WriteFile(target, data, info.Mode().Perm())
})
}
+70
View File
@@ -0,0 +1,70 @@
package goname
import (
"errors"
"io/fs"
"path/filepath"
"strings"
)
// ImportResult captures changes made by RenameImport.
type ImportResult struct {
OldImportPrefix string
NewImportPrefix string
FilesUpdated int
ImportsUpdated int
}
// RenameImport rewrites import paths that match oldImportPrefix within the directory tree.
func RenameImport(dir, oldImportPrefix, newImportPrefix string) (*ImportResult, error) {
oldImportPrefix = strings.TrimSpace(oldImportPrefix)
newImportPrefix = strings.TrimSpace(newImportPrefix)
if oldImportPrefix == "" {
return nil, errors.New("old import prefix is required")
}
if newImportPrefix == "" {
return nil, errors.New("new import prefix is required")
}
result := &ImportResult{
OldImportPrefix: oldImportPrefix,
NewImportPrefix: newImportPrefix,
}
if err := rewriteImportPrefixes(dir, oldImportPrefix, newImportPrefix, result); err != nil {
return result, err
}
if result.ImportsUpdated == 0 {
return result, errors.New("no imports updated (check old/new import prefixes)")
}
return result, nil
}
func rewriteImportPrefixes(dir, oldImportPrefix, newImportPrefix string, result *ImportResult) error {
return filepath.WalkDir(dir, func(path string, entry fs.DirEntry, err error) error {
if err != nil {
return err
}
if entry.IsDir() {
name := entry.Name()
if name == "vendor" || strings.HasPrefix(name, ".") {
return fs.SkipDir
}
return nil
}
if filepath.Ext(path) != ".go" {
return nil
}
changed, importsUpdated, err := rewriteFileImports(path, oldImportPrefix, newImportPrefix)
if err != nil {
return err
}
if changed {
result.FilesUpdated++
result.ImportsUpdated += importsUpdated
}
return nil
})
}
+91
View File
@@ -0,0 +1,91 @@
package goname
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestRenameImportBasic(t *testing.T) {
fixture := filepath.Join("testdata", "rename-import", "basic")
tmp := t.TempDir()
if err := copyDir(fixture, tmp); err != nil {
t.Fatalf("copy fixture: %v", err)
}
result, err := RenameImport(tmp, "foo/bar", "acme/corp")
if err != nil {
t.Fatalf("RenameImport: %v", err)
}
if result.OldImportPrefix != "foo/bar" {
t.Fatalf("old prefix mismatch: %s", result.OldImportPrefix)
}
if result.NewImportPrefix != "acme/corp" {
t.Fatalf("new prefix mismatch: %s", result.NewImportPrefix)
}
if result.FilesUpdated != 2 {
t.Fatalf("files updated mismatch: %d", result.FilesUpdated)
}
if result.ImportsUpdated != 4 {
t.Fatalf("imports updated mismatch: %d", result.ImportsUpdated)
}
mainImports := importPaths(t, filepath.Join(tmp, "main.go"))
for _, want := range []string{"acme/corp", "acme/corp/baz", "acme/corp/qux", "foo/barista", "other/pkg"} {
if !mainImports[want] {
t.Fatalf("main.go missing import: %s", want)
}
}
if mainImports["foo/bar"] || mainImports["foo/bar/baz"] || mainImports["foo/bar/qux"] {
t.Fatalf("main.go still references old prefix")
}
barImports := importPaths(t, filepath.Join(tmp, "foo", "bar", "bar.go"))
if !barImports["acme/corp/baz"] {
t.Fatalf("bar.go import not updated")
}
vendorImports := importPaths(t, filepath.Join(tmp, "vendor", "foo", "bar", "vendor.go"))
if !vendorImports["foo/bar/baz"] {
t.Fatalf("vendor import should remain unchanged")
}
hiddenImports := importPaths(t, filepath.Join(tmp, ".hidden", "hidden.go"))
if !hiddenImports["foo/bar/baz"] {
t.Fatalf("hidden import should remain unchanged")
}
mainBytes, err := os.ReadFile(filepath.Join(tmp, "main.go"))
if err != nil {
t.Fatalf("read main.go: %v", err)
}
if !strings.Contains(string(mainBytes), "const example = \"foo/bar/baz\"") {
t.Fatalf("string literal should remain unchanged")
}
}
func TestRenameImportNoChanges(t *testing.T) {
fixture := filepath.Join("testdata", "rename-import", "basic")
tmp := t.TempDir()
if err := copyDir(fixture, tmp); err != nil {
t.Fatalf("copy fixture: %v", err)
}
_, err := RenameImport(tmp, "missing/prefix", "acme/corp")
if err == nil {
t.Fatalf("expected error for missing prefix")
}
}
func TestRenameImportInvalidArgs(t *testing.T) {
_, err := RenameImport(t.TempDir(), "", "acme/corp")
if err == nil {
t.Fatalf("expected error for empty old prefix")
}
_, err = RenameImport(t.TempDir(), "foo/bar", "")
if err == nil {
t.Fatalf("expected error for empty new prefix")
}
}
+5
View File
@@ -0,0 +1,5 @@
module example.com/old
go 1.22
require example.com/other v1.2.3
+14
View File
@@ -0,0 +1,14 @@
package main
import (
"fmt"
"example.com/old"
oldpkg "example.com/old/pkg/foo"
_ "example.com/oldish/pkg"
)
const example = "example.com/old/pkg/foo"
func main() {
fmt.Println(oldpkg.Name, example)
}
+3
View File
@@ -0,0 +1,3 @@
package bar
const Name = "bar"
+5
View File
@@ -0,0 +1,5 @@
package foo
import "example.com/old/pkg/bar"
const Name = bar.Name
+5
View File
@@ -0,0 +1,5 @@
package vendor
import "example.com/old/pkg/foo"
const Vendor = foo.Name
@@ -0,0 +1,3 @@
package hidden
import _ "foo/bar/baz"
@@ -0,0 +1,9 @@
package bar
import "foo/bar/baz"
const Name = "bar"
func Use() string {
return baz.Name
}
@@ -0,0 +1,3 @@
package baz
const Name = "baz"
@@ -0,0 +1,3 @@
package qux
const Name = "qux"
@@ -0,0 +1,3 @@
package barista
const Name = "barista"
@@ -0,0 +1,3 @@
module example.com/demo
go 1.22
+16
View File
@@ -0,0 +1,16 @@
package main
import (
"fmt"
bar "foo/bar"
_ "foo/bar/baz"
_ "foo/bar/qux"
_ "foo/barista"
"other/pkg"
)
const example = "foo/bar/baz"
func main() {
fmt.Println(bar.Name, example)
}
@@ -0,0 +1,3 @@
package vendor
import _ "foo/bar/baz"
+27
View File
@@ -0,0 +1,27 @@
Copyright (c) 2009 The Go Authors. All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:
* Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above
copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the
distribution.
* Neither the name of Google Inc. nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+5
View File
@@ -0,0 +1,5 @@
# generic-list-go [![CI](https://github.com/bahlo/generic-list-go/actions/workflows/ci.yml/badge.svg)](https://github.com/bahlo/generic-list-go/actions/workflows/ci.yml)
Go [container/list](https://pkg.go.dev/container/list) but with generics.
The code is based on `container/list` in `go1.18beta2`.
+235
View File
@@ -0,0 +1,235 @@
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package list implements a doubly linked list.
//
// To iterate over a list (where l is a *List):
// for e := l.Front(); e != nil; e = e.Next() {
// // do something with e.Value
// }
//
package list
// Element is an element of a linked list.
type Element[T any] struct {
// Next and previous pointers in the doubly-linked list of elements.
// To simplify the implementation, internally a list l is implemented
// as a ring, such that &l.root is both the next element of the last
// list element (l.Back()) and the previous element of the first list
// element (l.Front()).
next, prev *Element[T]
// The list to which this element belongs.
list *List[T]
// The value stored with this element.
Value T
}
// Next returns the next list element or nil.
func (e *Element[T]) Next() *Element[T] {
if p := e.next; e.list != nil && p != &e.list.root {
return p
}
return nil
}
// Prev returns the previous list element or nil.
func (e *Element[T]) Prev() *Element[T] {
if p := e.prev; e.list != nil && p != &e.list.root {
return p
}
return nil
}
// List represents a doubly linked list.
// The zero value for List is an empty list ready to use.
type List[T any] struct {
root Element[T] // sentinel list element, only &root, root.prev, and root.next are used
len int // current list length excluding (this) sentinel element
}
// Init initializes or clears list l.
func (l *List[T]) Init() *List[T] {
l.root.next = &l.root
l.root.prev = &l.root
l.len = 0
return l
}
// New returns an initialized list.
func New[T any]() *List[T] { return new(List[T]).Init() }
// Len returns the number of elements of list l.
// The complexity is O(1).
func (l *List[T]) Len() int { return l.len }
// Front returns the first element of list l or nil if the list is empty.
func (l *List[T]) Front() *Element[T] {
if l.len == 0 {
return nil
}
return l.root.next
}
// Back returns the last element of list l or nil if the list is empty.
func (l *List[T]) Back() *Element[T] {
if l.len == 0 {
return nil
}
return l.root.prev
}
// lazyInit lazily initializes a zero List value.
func (l *List[T]) lazyInit() {
if l.root.next == nil {
l.Init()
}
}
// insert inserts e after at, increments l.len, and returns e.
func (l *List[T]) insert(e, at *Element[T]) *Element[T] {
e.prev = at
e.next = at.next
e.prev.next = e
e.next.prev = e
e.list = l
l.len++
return e
}
// insertValue is a convenience wrapper for insert(&Element{Value: v}, at).
func (l *List[T]) insertValue(v T, at *Element[T]) *Element[T] {
return l.insert(&Element[T]{Value: v}, at)
}
// remove removes e from its list, decrements l.len
func (l *List[T]) remove(e *Element[T]) {
e.prev.next = e.next
e.next.prev = e.prev
e.next = nil // avoid memory leaks
e.prev = nil // avoid memory leaks
e.list = nil
l.len--
}
// move moves e to next to at.
func (l *List[T]) move(e, at *Element[T]) {
if e == at {
return
}
e.prev.next = e.next
e.next.prev = e.prev
e.prev = at
e.next = at.next
e.prev.next = e
e.next.prev = e
}
// Remove removes e from l if e is an element of list l.
// It returns the element value e.Value.
// The element must not be nil.
func (l *List[T]) Remove(e *Element[T]) T {
if e.list == l {
// if e.list == l, l must have been initialized when e was inserted
// in l or l == nil (e is a zero Element) and l.remove will crash
l.remove(e)
}
return e.Value
}
// PushFront inserts a new element e with value v at the front of list l and returns e.
func (l *List[T]) PushFront(v T) *Element[T] {
l.lazyInit()
return l.insertValue(v, &l.root)
}
// PushBack inserts a new element e with value v at the back of list l and returns e.
func (l *List[T]) PushBack(v T) *Element[T] {
l.lazyInit()
return l.insertValue(v, l.root.prev)
}
// InsertBefore inserts a new element e with value v immediately before mark and returns e.
// If mark is not an element of l, the list is not modified.
// The mark must not be nil.
func (l *List[T]) InsertBefore(v T, mark *Element[T]) *Element[T] {
if mark.list != l {
return nil
}
// see comment in List.Remove about initialization of l
return l.insertValue(v, mark.prev)
}
// InsertAfter inserts a new element e with value v immediately after mark and returns e.
// If mark is not an element of l, the list is not modified.
// The mark must not be nil.
func (l *List[T]) InsertAfter(v T, mark *Element[T]) *Element[T] {
if mark.list != l {
return nil
}
// see comment in List.Remove about initialization of l
return l.insertValue(v, mark)
}
// MoveToFront moves element e to the front of list l.
// If e is not an element of l, the list is not modified.
// The element must not be nil.
func (l *List[T]) MoveToFront(e *Element[T]) {
if e.list != l || l.root.next == e {
return
}
// see comment in List.Remove about initialization of l
l.move(e, &l.root)
}
// MoveToBack moves element e to the back of list l.
// If e is not an element of l, the list is not modified.
// The element must not be nil.
func (l *List[T]) MoveToBack(e *Element[T]) {
if e.list != l || l.root.prev == e {
return
}
// see comment in List.Remove about initialization of l
l.move(e, l.root.prev)
}
// MoveBefore moves element e to its new position before mark.
// If e or mark is not an element of l, or e == mark, the list is not modified.
// The element and mark must not be nil.
func (l *List[T]) MoveBefore(e, mark *Element[T]) {
if e.list != l || e == mark || mark.list != l {
return
}
l.move(e, mark.prev)
}
// MoveAfter moves element e to its new position after mark.
// If e or mark is not an element of l, or e == mark, the list is not modified.
// The element and mark must not be nil.
func (l *List[T]) MoveAfter(e, mark *Element[T]) {
if e.list != l || e == mark || mark.list != l {
return
}
l.move(e, mark)
}
// PushBackList inserts a copy of another list at the back of list l.
// The lists l and other may be the same. They must not be nil.
func (l *List[T]) PushBackList(other *List[T]) {
l.lazyInit()
for i, e := other.Len(), other.Front(); i > 0; i, e = i-1, e.Next() {
l.insertValue(e.Value, l.root.prev)
}
}
// PushFrontList inserts a copy of another list at the front of list l.
// The lists l and other may be the same. They must not be nil.
func (l *List[T]) PushFrontList(other *List[T]) {
l.lazyInit()
for i, e := other.Len(), other.Back(); i > 0; i, e = i-1, e.Prev() {
l.insertValue(e.Value, &l.root)
}
}
+12
View File
@@ -0,0 +1,12 @@
*.test
*.out
*.mprof
.idea
vendor/github.com/buger/goterm/
prof.cpu
prof.mem
+11
View File
@@ -0,0 +1,11 @@
language: go
arch:
- amd64
- ppc64le
go:
- 1.7.x
- 1.8.x
- 1.9.x
- 1.10.x
- 1.11.x
script: go test -v ./.
+12
View File
@@ -0,0 +1,12 @@
FROM golang:1.6
RUN go get github.com/Jeffail/gabs
RUN go get github.com/bitly/go-simplejson
RUN go get github.com/pquerna/ffjson
RUN go get github.com/antonholmquist/jason
RUN go get github.com/mreiferson/go-ujson
RUN go get -tags=unsafe -u github.com/ugorji/go/codec
RUN go get github.com/mailru/easyjson
WORKDIR /go/src/github.com/buger/jsonparser
ADD . /go/src/github.com/buger/jsonparser
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2016 Leonid Bugaev
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+36
View File
@@ -0,0 +1,36 @@
SOURCE = parser.go
CONTAINER = jsonparser
SOURCE_PATH = /go/src/github.com/buger/jsonparser
BENCHMARK = JsonParser
BENCHTIME = 5s
TEST = .
DRUN = docker run -v `pwd`:$(SOURCE_PATH) -i -t $(CONTAINER)
build:
docker build -t $(CONTAINER) .
race:
$(DRUN) --env GORACE="halt_on_error=1" go test ./. $(ARGS) -v -race -timeout 15s
bench:
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -benchtime $(BENCHTIME) -v
bench_local:
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench . $(ARGS) -benchtime $(BENCHTIME) -v
profile:
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -memprofile mem.mprof -v
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -cpuprofile cpu.out -v
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -c
test:
$(DRUN) go test $(LDFLAGS) ./ -run $(TEST) -timeout 10s $(ARGS) -v
fmt:
$(DRUN) go fmt ./...
vet:
$(DRUN) go vet ./.
bash:
$(DRUN) /bin/bash
+365
View File
@@ -0,0 +1,365 @@
[![Go Report Card](https://goreportcard.com/badge/github.com/buger/jsonparser)](https://goreportcard.com/report/github.com/buger/jsonparser) ![License](https://img.shields.io/dub/l/vibe-d.svg)
# Alternative JSON parser for Go (10x times faster standard library)
It does not require you to know the structure of the payload (eg. create structs), and allows accessing fields by providing the path to them. It is up to **10 times faster** than standard `encoding/json` package (depending on payload size and usage), **allocates no memory**. See benchmarks below.
## Rationale
Originally I made this for a project that relies on a lot of 3rd party APIs that can be unpredictable and complex.
I love simplicity and prefer to avoid external dependecies. `encoding/json` requires you to know exactly your data structures, or if you prefer to use `map[string]interface{}` instead, it will be very slow and hard to manage.
I investigated what's on the market and found that most libraries are just wrappers around `encoding/json`, there is few options with own parsers (`ffjson`, `easyjson`), but they still requires you to create data structures.
Goal of this project is to push JSON parser to the performance limits and not sacrifice with compliance and developer user experience.
## Example
For the given JSON our goal is to extract the user's full name, number of github followers and avatar.
```go
import "github.com/buger/jsonparser"
...
data := []byte(`{
"person": {
"name": {
"first": "Leonid",
"last": "Bugaev",
"fullName": "Leonid Bugaev"
},
"github": {
"handle": "buger",
"followers": 109
},
"avatars": [
{ "url": "https://avatars1.githubusercontent.com/u/14009?v=3&s=460", "type": "thumbnail" }
]
},
"company": {
"name": "Acme"
}
}`)
// You can specify key path by providing arguments to Get function
jsonparser.Get(data, "person", "name", "fullName")
// There is `GetInt` and `GetBoolean` helpers if you exactly know key data type
jsonparser.GetInt(data, "person", "github", "followers")
// When you try to get object, it will return you []byte slice pointer to data containing it
// In `company` it will be `{"name": "Acme"}`
jsonparser.Get(data, "company")
// If the key doesn't exist it will throw an error
var size int64
if value, err := jsonparser.GetInt(data, "company", "size"); err == nil {
size = value
}
// You can use `ArrayEach` helper to iterate items [item1, item2 .... itemN]
jsonparser.ArrayEach(data, func(value []byte, dataType jsonparser.ValueType, offset int, err error) {
fmt.Println(jsonparser.Get(value, "url"))
}, "person", "avatars")
// Or use can access fields by index!
jsonparser.GetString(data, "person", "avatars", "[0]", "url")
// You can use `ObjectEach` helper to iterate objects { "key1":object1, "key2":object2, .... "keyN":objectN }
jsonparser.ObjectEach(data, func(key []byte, value []byte, dataType jsonparser.ValueType, offset int) error {
fmt.Printf("Key: '%s'\n Value: '%s'\n Type: %s\n", string(key), string(value), dataType)
return nil
}, "person", "name")
// The most efficient way to extract multiple keys is `EachKey`
paths := [][]string{
[]string{"person", "name", "fullName"},
[]string{"person", "avatars", "[0]", "url"},
[]string{"company", "url"},
}
jsonparser.EachKey(data, func(idx int, value []byte, vt jsonparser.ValueType, err error){
switch idx {
case 0: // []string{"person", "name", "fullName"}
...
case 1: // []string{"person", "avatars", "[0]", "url"}
...
case 2: // []string{"company", "url"},
...
}
}, paths...)
// For more information see docs below
```
## Need to speedup your app?
I'm available for consulting and can help you push your app performance to the limits. Ping me at: leonsbox@gmail.com.
## Reference
Library API is really simple. You just need the `Get` method to perform any operation. The rest is just helpers around it.
You also can view API at [godoc.org](https://godoc.org/github.com/buger/jsonparser)
### **`Get`**
```go
func Get(data []byte, keys ...string) (value []byte, dataType jsonparser.ValueType, offset int, err error)
```
Receives data structure, and key path to extract value from.
Returns:
* `value` - Pointer to original data structure containing key value, or just empty slice if nothing found or error
* `dataType` - Can be: `NotExist`, `String`, `Number`, `Object`, `Array`, `Boolean` or `Null`
* `offset` - Offset from provided data structure where key value ends. Used mostly internally, for example for `ArrayEach` helper.
* `err` - If the key is not found or any other parsing issue, it should return error. If key not found it also sets `dataType` to `NotExist`
Accepts multiple keys to specify path to JSON value (in case of quering nested structures).
If no keys are provided it will try to extract the closest JSON value (simple ones or object/array), useful for reading streams or arrays, see `ArrayEach` implementation.
Note that keys can be an array indexes: `jsonparser.GetInt("person", "avatars", "[0]", "url")`, pretty cool, yeah?
### **`GetString`**
```go
func GetString(data []byte, keys ...string) (val string, err error)
```
Returns strings properly handing escaped and unicode characters. Note that this will cause additional memory allocations.
### **`GetUnsafeString`**
If you need string in your app, and ready to sacrifice with support of escaped symbols in favor of speed. It returns string mapped to existing byte slice memory, without any allocations:
```go
s, _, := jsonparser.GetUnsafeString(data, "person", "name", "title")
switch s {
case 'CEO':
...
case 'Engineer'
...
...
}
```
Note that `unsafe` here means that your string will exist until GC will free underlying byte slice, for most of cases it means that you can use this string only in current context, and should not pass it anywhere externally: through channels or any other way.
### **`GetBoolean`**, **`GetInt`** and **`GetFloat`**
```go
func GetBoolean(data []byte, keys ...string) (val bool, err error)
func GetFloat(data []byte, keys ...string) (val float64, err error)
func GetInt(data []byte, keys ...string) (val int64, err error)
```
If you know the key type, you can use the helpers above.
If key data type do not match, it will return error.
### **`ArrayEach`**
```go
func ArrayEach(data []byte, cb func(value []byte, dataType jsonparser.ValueType, offset int, err error), keys ...string)
```
Needed for iterating arrays, accepts a callback function with the same return arguments as `Get`.
### **`ObjectEach`**
```go
func ObjectEach(data []byte, callback func(key []byte, value []byte, dataType ValueType, offset int) error, keys ...string) (err error)
```
Needed for iterating object, accepts a callback function. Example:
```go
var handler func([]byte, []byte, jsonparser.ValueType, int) error
handler = func(key []byte, value []byte, dataType jsonparser.ValueType, offset int) error {
//do stuff here
}
jsonparser.ObjectEach(myJson, handler)
```
### **`EachKey`**
```go
func EachKey(data []byte, cb func(idx int, value []byte, dataType jsonparser.ValueType, err error), paths ...[]string)
```
When you need to read multiple keys, and you do not afraid of low-level API `EachKey` is your friend. It read payload only single time, and calls callback function once path is found. For example when you call multiple times `Get`, it has to process payload multiple times, each time you call it. Depending on payload `EachKey` can be multiple times faster than `Get`. Path can use nested keys as well!
```go
paths := [][]string{
[]string{"uuid"},
[]string{"tz"},
[]string{"ua"},
[]string{"st"},
}
var data SmallPayload
jsonparser.EachKey(smallFixture, func(idx int, value []byte, vt jsonparser.ValueType, err error){
switch idx {
case 0:
data.Uuid, _ = value
case 1:
v, _ := jsonparser.ParseInt(value)
data.Tz = int(v)
case 2:
data.Ua, _ = value
case 3:
v, _ := jsonparser.ParseInt(value)
data.St = int(v)
}
}, paths...)
```
### **`Set`**
```go
func Set(data []byte, setValue []byte, keys ...string) (value []byte, err error)
```
Receives existing data structure, key path to set, and value to set at that key. *This functionality is experimental.*
Returns:
* `value` - Pointer to original data structure with updated or added key value.
* `err` - If any parsing issue, it should return error.
Accepts multiple keys to specify path to JSON value (in case of updating or creating nested structures).
Note that keys can be an array indexes: `jsonparser.Set(data, []byte("http://github.com"), "person", "avatars", "[0]", "url")`
### **`Delete`**
```go
func Delete(data []byte, keys ...string) value []byte
```
Receives existing data structure, and key path to delete. *This functionality is experimental.*
Returns:
* `value` - Pointer to original data structure with key path deleted if it can be found. If there is no key path, then the whole data structure is deleted.
Accepts multiple keys to specify path to JSON value (in case of updating or creating nested structures).
Note that keys can be an array indexes: `jsonparser.Delete(data, "person", "avatars", "[0]", "url")`
## What makes it so fast?
* It does not rely on `encoding/json`, `reflection` or `interface{}`, the only real package dependency is `bytes`.
* Operates with JSON payload on byte level, providing you pointers to the original data structure: no memory allocation.
* No automatic type conversions, by default everything is a []byte, but it provides you value type, so you can convert by yourself (there is few helpers included).
* Does not parse full record, only keys you specified
## Benchmarks
There are 3 benchmark types, trying to simulate real-life usage for small, medium and large JSON payloads.
For each metric, the lower value is better. Time/op is in nanoseconds. Values better than standard encoding/json marked as bold text.
Benchmarks run on standard Linode 1024 box.
Compared libraries:
* https://golang.org/pkg/encoding/json
* https://github.com/Jeffail/gabs
* https://github.com/a8m/djson
* https://github.com/bitly/go-simplejson
* https://github.com/antonholmquist/jason
* https://github.com/mreiferson/go-ujson
* https://github.com/ugorji/go/codec
* https://github.com/pquerna/ffjson
* https://github.com/mailru/easyjson
* https://github.com/buger/jsonparser
#### TLDR
If you want to skip next sections we have 2 winner: `jsonparser` and `easyjson`.
`jsonparser` is up to 10 times faster than standard `encoding/json` package (depending on payload size and usage), and almost infinitely (literally) better in memory consumption because it operates with data on byte level, and provide direct slice pointers.
`easyjson` wins in CPU in medium tests and frankly i'm impressed with this package: it is remarkable results considering that it is almost drop-in replacement for `encoding/json` (require some code generation).
It's hard to fully compare `jsonparser` and `easyjson` (or `ffson`), they a true parsers and fully process record, unlike `jsonparser` which parse only keys you specified.
If you searching for replacement of `encoding/json` while keeping structs, `easyjson` is an amazing choice. If you want to process dynamic JSON, have memory constrains, or more control over your data you should try `jsonparser`.
`jsonparser` performance heavily depends on usage, and it works best when you do not need to process full record, only some keys. The more calls you need to make, the slower it will be, in contrast `easyjson` (or `ffjson`, `encoding/json`) parser record only 1 time, and then you can make as many calls as you want.
With great power comes great responsibility! :)
#### Small payload
Each test processes 190 bytes of http log as a JSON record.
It should read multiple fields.
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_small_payload_test.go
Library | time/op | bytes/op | allocs/op
------ | ------- | -------- | -------
encoding/json struct | 7879 | 880 | 18
encoding/json interface{} | 8946 | 1521 | 38
Jeffail/gabs | 10053 | 1649 | 46
bitly/go-simplejson | 10128 | 2241 | 36
antonholmquist/jason | 27152 | 7237 | 101
github.com/ugorji/go/codec | 8806 | 2176 | 31
mreiferson/go-ujson | **7008** | **1409** | 37
a8m/djson | 3862 | 1249 | 30
pquerna/ffjson | **3769** | **624** | **15**
mailru/easyjson | **2002** | **192** | **9**
buger/jsonparser | **1367** | **0** | **0**
buger/jsonparser (EachKey API) | **809** | **0** | **0**
Winners are ffjson, easyjson and jsonparser, where jsonparser is up to 9.8x faster than encoding/json and 4.6x faster than ffjson, and slightly faster than easyjson.
If you look at memory allocation, jsonparser has no rivals, as it makes no data copy and operates with raw []byte structures and pointers to it.
#### Medium payload
Each test processes a 2.4kb JSON record (based on Clearbit API).
It should read multiple nested fields and 1 array.
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_medium_payload_test.go
| Library | time/op | bytes/op | allocs/op |
| ------- | ------- | -------- | --------- |
| encoding/json struct | 57749 | 1336 | 29 |
| encoding/json interface{} | 79297 | 10627 | 215 |
| Jeffail/gabs | 83807 | 11202 | 235 |
| bitly/go-simplejson | 88187 | 17187 | 220 |
| antonholmquist/jason | 94099 | 19013 | 247 |
| github.com/ugorji/go/codec | 114719 | 6712 | 152 |
| mreiferson/go-ujson | **56972** | 11547 | 270 |
| a8m/djson | 28525 | 10196 | 198 |
| pquerna/ffjson | **20298** | **856** | **20** |
| mailru/easyjson | **10512** | **336** | **12** |
| buger/jsonparser | **15955** | **0** | **0** |
| buger/jsonparser (EachKey API) | **8916** | **0** | **0** |
The difference between ffjson and jsonparser in CPU usage is smaller, while the memory consumption difference is growing. On the other hand `easyjson` shows remarkable performance for medium payload.
`gabs`, `go-simplejson` and `jason` are based on encoding/json and map[string]interface{} and actually only helpers for unstructured JSON, their performance correlate with `encoding/json interface{}`, and they will skip next round.
`go-ujson` while have its own parser, shows same performance as `encoding/json`, also skips next round. Same situation with `ugorji/go/codec`, but it showed unexpectedly bad performance for complex payloads.
#### Large payload
Each test processes a 24kb JSON record (based on Discourse API)
It should read 2 arrays, and for each item in array get a few fields.
Basically it means processing a full JSON file.
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_large_payload_test.go
| Library | time/op | bytes/op | allocs/op |
| --- | --- | --- | --- |
| encoding/json struct | 748336 | 8272 | 307 |
| encoding/json interface{} | 1224271 | 215425 | 3395 |
| a8m/djson | 510082 | 213682 | 2845 |
| pquerna/ffjson | **312271** | **7792** | **298** |
| mailru/easyjson | **154186** | **6992** | **288** |
| buger/jsonparser | **85308** | **0** | **0** |
`jsonparser` now is a winner, but do not forget that it is way more lightweight parser than `ffson` or `easyjson`, and they have to parser all the data, while `jsonparser` parse only what you need. All `ffjson`, `easysjon` and `jsonparser` have their own parsing code, and does not depend on `encoding/json` or `interface{}`, thats one of the reasons why they are so fast. `easyjson` also use a bit of `unsafe` package to reduce memory consuption (in theory it can lead to some unexpected GC issue, but i did not tested enough)
Also last benchmark did not included `EachKey` test, because in this particular case we need to read lot of Array values, and using `ArrayEach` is more efficient.
## Questions and support
All bug-reports and suggestions should go though Github Issues.
## Contributing
1. Fork it
2. Create your feature branch (git checkout -b my-new-feature)
3. Commit your changes (git commit -am 'Added some feature')
4. Push to the branch (git push origin my-new-feature)
5. Create new Pull Request
## Development
All my development happens using Docker, and repo include some Make tasks to simplify development.
* `make build` - builds docker image, usually can be called only once
* `make test` - run tests
* `make fmt` - run go fmt
* `make bench` - run benchmarks (if you need to run only single benchmark modify `BENCHMARK` variable in make file)
* `make profile` - runs benchmark and generate 3 files- `cpu.out`, `mem.mprof` and `benchmark.test` binary, which can be used for `go tool pprof`
* `make bash` - enter container (i use it for running `go tool pprof` above)
+47
View File
@@ -0,0 +1,47 @@
package jsonparser
import (
bio "bytes"
)
// minInt64 '-9223372036854775808' is the smallest representable number in int64
const minInt64 = `9223372036854775808`
// About 2x faster then strconv.ParseInt because it only supports base 10, which is enough for JSON
func parseInt(bytes []byte) (v int64, ok bool, overflow bool) {
if len(bytes) == 0 {
return 0, false, false
}
var neg bool = false
if bytes[0] == '-' {
neg = true
bytes = bytes[1:]
}
var b int64 = 0
for _, c := range bytes {
if c >= '0' && c <= '9' {
b = (10 * v) + int64(c-'0')
} else {
return 0, false, false
}
if overflow = (b < v); overflow {
break
}
v = b
}
if overflow {
if neg && bio.Equal(bytes, []byte(minInt64)) {
return b, true, false
}
return 0, false, true
}
if neg {
return -v, true, false
} else {
return v, true, false
}
}
+25
View File
@@ -0,0 +1,25 @@
// +build appengine appenginevm
package jsonparser
import (
"strconv"
)
// See fastbytes_unsafe.go for explanation on why *[]byte is used (signatures must be consistent with those in that file)
func equalStr(b *[]byte, s string) bool {
return string(*b) == s
}
func parseFloat(b *[]byte) (float64, error) {
return strconv.ParseFloat(string(*b), 64)
}
func bytesToString(b *[]byte) string {
return string(*b)
}
func StringToBytes(s string) []byte {
return []byte(s)
}
+44
View File
@@ -0,0 +1,44 @@
// +build !appengine,!appenginevm
package jsonparser
import (
"reflect"
"strconv"
"unsafe"
"runtime"
)
//
// The reason for using *[]byte rather than []byte in parameters is an optimization. As of Go 1.6,
// the compiler cannot perfectly inline the function when using a non-pointer slice. That is,
// the non-pointer []byte parameter version is slower than if its function body is manually
// inlined, whereas the pointer []byte version is equally fast to the manually inlined
// version. Instruction count in assembly taken from "go tool compile" confirms this difference.
//
// TODO: Remove hack after Go 1.7 release
//
func equalStr(b *[]byte, s string) bool {
return *(*string)(unsafe.Pointer(b)) == s
}
func parseFloat(b *[]byte) (float64, error) {
return strconv.ParseFloat(*(*string)(unsafe.Pointer(b)), 64)
}
// A hack until issue golang/go#2632 is fixed.
// See: https://github.com/golang/go/issues/2632
func bytesToString(b *[]byte) string {
return *(*string)(unsafe.Pointer(b))
}
func StringToBytes(s string) []byte {
b := make([]byte, 0, 0)
bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
sh := (*reflect.StringHeader)(unsafe.Pointer(&s))
bh.Data = sh.Data
bh.Cap = sh.Len
bh.Len = sh.Len
runtime.KeepAlive(s)
return b
}
+173
View File
@@ -0,0 +1,173 @@
package jsonparser
import (
"bytes"
"unicode/utf8"
)
// JSON Unicode stuff: see https://tools.ietf.org/html/rfc7159#section-7
const supplementalPlanesOffset = 0x10000
const highSurrogateOffset = 0xD800
const lowSurrogateOffset = 0xDC00
const basicMultilingualPlaneReservedOffset = 0xDFFF
const basicMultilingualPlaneOffset = 0xFFFF
func combineUTF16Surrogates(high, low rune) rune {
return supplementalPlanesOffset + (high-highSurrogateOffset)<<10 + (low - lowSurrogateOffset)
}
const badHex = -1
func h2I(c byte) int {
switch {
case c >= '0' && c <= '9':
return int(c - '0')
case c >= 'A' && c <= 'F':
return int(c - 'A' + 10)
case c >= 'a' && c <= 'f':
return int(c - 'a' + 10)
}
return badHex
}
// decodeSingleUnicodeEscape decodes a single \uXXXX escape sequence. The prefix \u is assumed to be present and
// is not checked.
// In JSON, these escapes can either come alone or as part of "UTF16 surrogate pairs" that must be handled together.
// This function only handles one; decodeUnicodeEscape handles this more complex case.
func decodeSingleUnicodeEscape(in []byte) (rune, bool) {
// We need at least 6 characters total
if len(in) < 6 {
return utf8.RuneError, false
}
// Convert hex to decimal
h1, h2, h3, h4 := h2I(in[2]), h2I(in[3]), h2I(in[4]), h2I(in[5])
if h1 == badHex || h2 == badHex || h3 == badHex || h4 == badHex {
return utf8.RuneError, false
}
// Compose the hex digits
return rune(h1<<12 + h2<<8 + h3<<4 + h4), true
}
// isUTF16EncodedRune checks if a rune is in the range for non-BMP characters,
// which is used to describe UTF16 chars.
// Source: https://en.wikipedia.org/wiki/Plane_(Unicode)#Basic_Multilingual_Plane
func isUTF16EncodedRune(r rune) bool {
return highSurrogateOffset <= r && r <= basicMultilingualPlaneReservedOffset
}
func decodeUnicodeEscape(in []byte) (rune, int) {
if r, ok := decodeSingleUnicodeEscape(in); !ok {
// Invalid Unicode escape
return utf8.RuneError, -1
} else if r <= basicMultilingualPlaneOffset && !isUTF16EncodedRune(r) {
// Valid Unicode escape in Basic Multilingual Plane
return r, 6
} else if r2, ok := decodeSingleUnicodeEscape(in[6:]); !ok { // Note: previous decodeSingleUnicodeEscape success guarantees at least 6 bytes remain
// UTF16 "high surrogate" without manditory valid following Unicode escape for the "low surrogate"
return utf8.RuneError, -1
} else if r2 < lowSurrogateOffset {
// Invalid UTF16 "low surrogate"
return utf8.RuneError, -1
} else {
// Valid UTF16 surrogate pair
return combineUTF16Surrogates(r, r2), 12
}
}
// backslashCharEscapeTable: when '\X' is found for some byte X, it is to be replaced with backslashCharEscapeTable[X]
var backslashCharEscapeTable = [...]byte{
'"': '"',
'\\': '\\',
'/': '/',
'b': '\b',
'f': '\f',
'n': '\n',
'r': '\r',
't': '\t',
}
// unescapeToUTF8 unescapes the single escape sequence starting at 'in' into 'out' and returns
// how many characters were consumed from 'in' and emitted into 'out'.
// If a valid escape sequence does not appear as a prefix of 'in', (-1, -1) to signal the error.
func unescapeToUTF8(in, out []byte) (inLen int, outLen int) {
if len(in) < 2 || in[0] != '\\' {
// Invalid escape due to insufficient characters for any escape or no initial backslash
return -1, -1
}
// https://tools.ietf.org/html/rfc7159#section-7
switch e := in[1]; e {
case '"', '\\', '/', 'b', 'f', 'n', 'r', 't':
// Valid basic 2-character escapes (use lookup table)
out[0] = backslashCharEscapeTable[e]
return 2, 1
case 'u':
// Unicode escape
if r, inLen := decodeUnicodeEscape(in); inLen == -1 {
// Invalid Unicode escape
return -1, -1
} else {
// Valid Unicode escape; re-encode as UTF8
outLen := utf8.EncodeRune(out, r)
return inLen, outLen
}
}
return -1, -1
}
// unescape unescapes the string contained in 'in' and returns it as a slice.
// If 'in' contains no escaped characters:
// Returns 'in'.
// Else, if 'out' is of sufficient capacity (guaranteed if cap(out) >= len(in)):
// 'out' is used to build the unescaped string and is returned with no extra allocation
// Else:
// A new slice is allocated and returned.
func Unescape(in, out []byte) ([]byte, error) {
firstBackslash := bytes.IndexByte(in, '\\')
if firstBackslash == -1 {
return in, nil
}
// Get a buffer of sufficient size (allocate if needed)
if cap(out) < len(in) {
out = make([]byte, len(in))
} else {
out = out[0:len(in)]
}
// Copy the first sequence of unescaped bytes to the output and obtain a buffer pointer (subslice)
copy(out, in[:firstBackslash])
in = in[firstBackslash:]
buf := out[firstBackslash:]
for len(in) > 0 {
// Unescape the next escaped character
inLen, bufLen := unescapeToUTF8(in, buf)
if inLen == -1 {
return nil, MalformedStringEscapeError
}
in = in[inLen:]
buf = buf[bufLen:]
// Copy everything up until the next backslash
nextBackslash := bytes.IndexByte(in, '\\')
if nextBackslash == -1 {
copy(buf, in)
buf = buf[len(in):]
break
} else {
copy(buf, in[:nextBackslash])
buf = buf[nextBackslash:]
in = in[nextBackslash:]
}
}
// Trim the out buffer to the amount that was actually emitted
return out[:len(out)-len(buf)], nil
}
+117
View File
@@ -0,0 +1,117 @@
package jsonparser
func FuzzParseString(data []byte) int {
r, err := ParseString(data)
if err != nil || r == "" {
return 0
}
return 1
}
func FuzzEachKey(data []byte) int {
paths := [][]string{
{"name"},
{"order"},
{"nested", "a"},
{"nested", "b"},
{"nested2", "a"},
{"nested", "nested3", "b"},
{"arr", "[1]", "b"},
{"arrInt", "[3]"},
{"arrInt", "[5]"},
{"nested"},
{"arr", "["},
{"a\n", "b\n"},
}
EachKey(data, func(idx int, value []byte, vt ValueType, err error) {}, paths...)
return 1
}
func FuzzDelete(data []byte) int {
Delete(data, "test")
return 1
}
func FuzzSet(data []byte) int {
_, err := Set(data, []byte(`"new value"`), "test")
if err != nil {
return 0
}
return 1
}
func FuzzObjectEach(data []byte) int {
_ = ObjectEach(data, func(key, value []byte, valueType ValueType, off int) error {
return nil
})
return 1
}
func FuzzParseFloat(data []byte) int {
_, err := ParseFloat(data)
if err != nil {
return 0
}
return 1
}
func FuzzParseInt(data []byte) int {
_, err := ParseInt(data)
if err != nil {
return 0
}
return 1
}
func FuzzParseBool(data []byte) int {
_, err := ParseBoolean(data)
if err != nil {
return 0
}
return 1
}
func FuzzTokenStart(data []byte) int {
_ = tokenStart(data)
return 1
}
func FuzzGetString(data []byte) int {
_, err := GetString(data, "test")
if err != nil {
return 0
}
return 1
}
func FuzzGetFloat(data []byte) int {
_, err := GetFloat(data, "test")
if err != nil {
return 0
}
return 1
}
func FuzzGetInt(data []byte) int {
_, err := GetInt(data, "test")
if err != nil {
return 0
}
return 1
}
func FuzzGetBoolean(data []byte) int {
_, err := GetBoolean(data, "test")
if err != nil {
return 0
}
return 1
}
func FuzzGetUnsafeString(data []byte) int {
_, err := GetUnsafeString(data, "test")
if err != nil {
return 0
}
return 1
}
+47
View File
@@ -0,0 +1,47 @@
#!/bin/bash -eu
git clone https://github.com/dvyukov/go-fuzz-corpus
zip corpus.zip go-fuzz-corpus/json/corpus/*
cp corpus.zip $OUT/fuzzparsestring_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzParseString fuzzparsestring
cp corpus.zip $OUT/fuzzeachkey_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzEachKey fuzzeachkey
cp corpus.zip $OUT/fuzzdelete_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzDelete fuzzdelete
cp corpus.zip $OUT/fuzzset_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzSet fuzzset
cp corpus.zip $OUT/fuzzobjecteach_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzObjectEach fuzzobjecteach
cp corpus.zip $OUT/fuzzparsefloat_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzParseFloat fuzzparsefloat
cp corpus.zip $OUT/fuzzparseint_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzParseInt fuzzparseint
cp corpus.zip $OUT/fuzzparsebool_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzParseBool fuzzparsebool
cp corpus.zip $OUT/fuzztokenstart_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzTokenStart fuzztokenstart
cp corpus.zip $OUT/fuzzgetstring_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzGetString fuzzgetstring
cp corpus.zip $OUT/fuzzgetfloat_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzGetFloat fuzzgetfloat
cp corpus.zip $OUT/fuzzgetint_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzGetInt fuzzgetint
cp corpus.zip $OUT/fuzzgetboolean_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzGetBoolean fuzzgetboolean
cp corpus.zip $OUT/fuzzgetunsafestring_seed_corpus.zip
compile_go_fuzzer github.com/buger/jsonparser FuzzGetUnsafeString fuzzgetunsafestring
+1283
View File
File diff suppressed because it is too large. Load diff
+2
View File
@@ -0,0 +1,2 @@
vendor/
.idea/
+69
View File
@@ -0,0 +1,69 @@
run:
tests: true
max-same-issues: 50
output:
print-issued-lines: false
linters:
enable:
- gocyclo
- gocritic
- goconst
- dupl
- unconvert
- goimports
- unused
- govet
- nakedret
- errcheck
- revive
- ineffassign
- goconst
- unparam
- gofmt
linters-settings:
vet:
check-shadowing: true
use-installed-packages: true
dupl:
threshold: 100
goconst:
min-len: 8
min-occurrences: 3
gocyclo:
min-complexity: 20
gocritic:
disabled-checks:
- ifElseChain
gofmt:
rewrite-rules:
- pattern: "interface{}"
replacement: "any"
- pattern: "a[b:len(a)]"
replacement: "a[b:]"
issues:
max-per-linter: 0
max-same: 0
exclude-dirs:
- resources
- old
exclude-files:
- cmd/protopkg/main.go
exclude-use-default: false
exclude:
# Captured by errcheck.
- "^(G104|G204):"
# Very commonly not checked.
- 'Error return value of .(.*\.Help|.*\.MarkFlagRequired|(os\.)?std(out|err)\..*|.*Close|.*Flush|os\.Remove(All)?|.*Print(f|ln|)|os\.(Un)?Setenv). is not checked'
# Weird error only seen on Kochiku...
- "internal error: no range for"
- 'exported method `.*\.(MarshalJSON|UnmarshalJSON|URN|Payload|GoString|Close|Provides|Requires|ExcludeFromHash|MarshalText|UnmarshalText|Description|Check|Poll|Severity)` should have comment or be unexported'
- "composite literal uses unkeyed fields"
- 'declaration of "err" shadows declaration'
- "by other packages, and that stutters"
- "Potential file inclusion via variable"
- "at least one file in a package should have a package comment"
- "bad syntax for struct tag pair"
+19
View File
@@ -0,0 +1,19 @@
Copyright (C) 2014 Alec Thomas
Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+374
View File
@@ -0,0 +1,374 @@
# Go JSON Schema Reflection
[![Lint](https://github.com/invopop/jsonschema/actions/workflows/lint.yaml/badge.svg)](https://github.com/invopop/jsonschema/actions/workflows/lint.yaml)
[![Test Go](https://github.com/invopop/jsonschema/actions/workflows/test.yaml/badge.svg)](https://github.com/invopop/jsonschema/actions/workflows/test.yaml)
[![Go Report Card](https://goreportcard.com/badge/github.com/invopop/jsonschema)](https://goreportcard.com/report/github.com/invopop/jsonschema)
[![GoDoc](https://godoc.org/github.com/invopop/jsonschema?status.svg)](https://godoc.org/github.com/invopop/jsonschema)
[![codecov](https://codecov.io/gh/invopop/jsonschema/graph/badge.svg?token=JMEB8W8GNZ)](https://codecov.io/gh/invopop/jsonschema)
![Latest Tag](https://img.shields.io/github/v/tag/invopop/jsonschema)
This package can be used to generate [JSON Schemas](http://json-schema.org/latest/json-schema-validation.html) from Go types through reflection.
- Supports arbitrarily complex types, including `interface{}`, maps, slices, etc.
- Supports json-schema features such as minLength, maxLength, pattern, format, etc.
- Supports simple string and numeric enums.
- Supports custom property fields via the `jsonschema_extras` struct tag.
This repository is a fork of the original [jsonschema](https://github.com/alecthomas/jsonschema) by [@alecthomas](https://github.com/alecthomas). At [Invopop](https://invopop.com) we use jsonschema as a cornerstone in our [GOBL library](https://github.com/invopop/gobl), and wanted to be able to continue building and adding features without taking up Alec's time. There have been a few significant changes that probably mean this version is a not compatible with with Alec's:
- The original was stuck on the draft-04 version of JSON Schema, we've now moved to the latest JSON Schema Draft 2020-12.
- Schema IDs are added automatically from the current Go package's URL in order to be unique, and can be disabled with the `Anonymous` option.
- Support for the `FullyQualifyTypeName` option has been removed. If you have conflicts, you should use multiple schema files with different IDs, set the `DoNotReference` option to true to hide definitions completely, or add your own naming strategy using the `Namer` property.
- Support for `yaml` tags and related options has been dropped for the sake of simplification. There were a [few inconsistencies](https://github.com/invopop/jsonschema/pull/21) around this that have now been fixed.
## Versions
This project is still under v0 scheme, as per Go convention, breaking changes are likely. Please pin go modules to version tags or branches, and reach out if you think something can be improved.
Go version >= 1.18 is required as generics are now being used.
## Example
The following Go type:
```go
type TestUser struct {
ID int `json:"id"`
Name string `json:"name" jsonschema:"title=the name,description=The name of a friend,example=joe,example=lucy,default=alex"`
Friends []int `json:"friends,omitempty" jsonschema_description:"The list of IDs, omitted when empty"`
Tags map[string]interface{} `json:"tags,omitempty" jsonschema_extras:"a=b,foo=bar,foo=bar1"`
BirthDate time.Time `json:"birth_date,omitempty" jsonschema:"oneof_required=date"`
YearOfBirth string `json:"year_of_birth,omitempty" jsonschema:"oneof_required=year"`
Metadata interface{} `json:"metadata,omitempty" jsonschema:"oneof_type=string;array"`
FavColor string `json:"fav_color,omitempty" jsonschema:"enum=red,enum=green,enum=blue"`
}
```
Results in following JSON Schema:
```go
jsonschema.Reflect(&TestUser{})
```
```json
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "https://github.com/invopop/jsonschema_test/test-user",
"$ref": "#/$defs/TestUser",
"$defs": {
"TestUser": {
"oneOf": [
{
"required": ["birth_date"],
"title": "date"
},
{
"required": ["year_of_birth"],
"title": "year"
}
],
"properties": {
"id": {
"type": "integer"
},
"name": {
"type": "string",
"title": "the name",
"description": "The name of a friend",
"default": "alex",
"examples": ["joe", "lucy"]
},
"friends": {
"items": {
"type": "integer"
},
"type": "array",
"description": "The list of IDs, omitted when empty"
},
"tags": {
"type": "object",
"a": "b",
"foo": ["bar", "bar1"]
},
"birth_date": {
"type": "string",
"format": "date-time"
},
"year_of_birth": {
"type": "string"
},
"metadata": {
"oneOf": [
{
"type": "string"
},
{
"type": "array"
}
]
},
"fav_color": {
"type": "string",
"enum": ["red", "green", "blue"]
}
},
"additionalProperties": false,
"type": "object",
"required": ["id", "name"]
}
}
}
```
## YAML
Support for `yaml` tags has now been removed. If you feel very strongly about this, we've opened a discussion to hear your comments: https://github.com/invopop/jsonschema/discussions/28
The recommended approach if you need to deal with YAML data is to first convert to JSON. The [invopop/yaml](https://github.com/invopop/yaml) library will make this trivial.
## Configurable behaviour
The behaviour of the schema generator can be altered with parameters when a `jsonschema.Reflector`
instance is created.
### ExpandedStruct
If set to `true`, makes the top level struct not to reference itself in the definitions. But type passed should be a struct type.
eg.
```go
type GrandfatherType struct {
FamilyName string `json:"family_name" jsonschema:"required"`
}
type SomeBaseType struct {
SomeBaseProperty int `json:"some_base_property"`
// The jsonschema required tag is nonsensical for private and ignored properties.
// Their presence here tests that the fields *will not* be required in the output
// schema, even if they are tagged required.
somePrivateBaseProperty string `json:"i_am_private" jsonschema:"required"`
SomeIgnoredBaseProperty string `json:"-" jsonschema:"required"`
SomeSchemaIgnoredProperty string `jsonschema:"-,required"`
SomeUntaggedBaseProperty bool `jsonschema:"required"`
someUnexportedUntaggedBaseProperty bool
Grandfather GrandfatherType `json:"grand"`
}
```
will output:
```json
{
"$schema": "http://json-schema.org/draft/2020-12/schema",
"required": ["some_base_property", "grand", "SomeUntaggedBaseProperty"],
"properties": {
"SomeUntaggedBaseProperty": {
"type": "boolean"
},
"grand": {
"$schema": "http://json-schema.org/draft/2020-12/schema",
"$ref": "#/definitions/GrandfatherType"
},
"some_base_property": {
"type": "integer"
}
},
"type": "object",
"$defs": {
"GrandfatherType": {
"required": ["family_name"],
"properties": {
"family_name": {
"type": "string"
}
},
"additionalProperties": false,
"type": "object"
}
}
}
```
### Using Go Comments
Writing a good schema with descriptions inside tags can become cumbersome and tedious, especially if you already have some Go comments around your types and field definitions. If you'd like to take advantage of these existing comments, you can use the `AddGoComments(base, path string)` method that forms part of the reflector to parse your go files and automatically generate a dictionary of Go import paths, types, and fields, to individual comments. These will then be used automatically as description fields, and can be overridden with a manual definition if needed.
Take a simplified example of a User struct which for the sake of simplicity we assume is defined inside this package:
```go
package main
// User is used as a base to provide tests for comments.
type User struct {
// Unique sequential identifier.
ID int `json:"id" jsonschema:"required"`
// Name of the user
Name string `json:"name"`
}
```
To get the comments provided into your JSON schema, use a regular `Reflector` and add the go code using an import module URL and path. Fully qualified go module paths cannot be determined reliably by the `go/parser` library, so we need to introduce this manually:
```go
r := new(Reflector)
if err := r.AddGoComments("github.com/invopop/jsonschema", "./"); err != nil {
// deal with error
}
s := r.Reflect(&User{})
// output
```
Expect the results to be similar to:
```json
{
"$schema": "http://json-schema.org/draft/2020-12/schema",
"$ref": "#/$defs/User",
"$defs": {
"User": {
"required": ["id"],
"properties": {
"id": {
"type": "integer",
"description": "Unique sequential identifier."
},
"name": {
"type": "string",
"description": "Name of the user"
}
},
"additionalProperties": false,
"type": "object",
"description": "User is used as a base to provide tests for comments."
}
}
}
```
### Custom Key Naming
In some situations, the keys actually used to write files are different from Go structs'.
This is often the case when writing a configuration file to YAML or JSON from a Go struct, or when returning a JSON response for a Web API: APIs typically use snake_case, while Go uses PascalCase.
You can pass a `func(string) string` function to `Reflector`'s `KeyNamer` option to map Go field names to JSON key names and reflect the aforementioned transformations, without having to specify `json:"..."` on every struct field.
For example, consider the following struct
```go
type User struct {
GivenName string
PasswordSalted []byte `json:"salted_password"`
}
```
We can transform field names to snake_case in the generated JSON schema:
```go
r := new(jsonschema.Reflector)
r.KeyNamer = strcase.SnakeCase // from package github.com/stoewer/go-strcase
r.Reflect(&User{})
```
Will yield
```diff
{
"$schema": "http://json-schema.org/draft/2020-12/schema",
"$ref": "#/$defs/User",
"$defs": {
"User": {
"properties": {
- "GivenName": {
+ "given_name": {
"type": "string"
},
"salted_password": {
"type": "string",
"contentEncoding": "base64"
}
},
"additionalProperties": false,
"type": "object",
- "required": ["GivenName", "salted_password"]
+ "required": ["given_name", "salted_password"]
}
}
}
```
As you can see, if a field name has a `json:""` tag set, the `key` argument to `KeyNamer` will have the value of that tag.
### Custom Type Definitions
Sometimes it can be useful to have custom JSON Marshal and Unmarshal methods in your structs that automatically convert for example a string into an object.
This library will recognize and attempt to call four different methods that help you adjust schemas to your specific needs:
- `JSONSchema() *Schema` - will prevent auto-generation of the schema so that you can provide your own definition.
- `JSONSchemaExtend(schema *jsonschema.Schema)` - will be called _after_ the schema has been generated, allowing you to add or manipulate the fields easily.
- `JSONSchemaAlias() any` - is called when reflecting the type of object and allows for an alternative to be used instead.
- `JSONSchemaProperty(prop string) any` - will be called for every property inside a struct giving you the chance to provide an alternative object to convert into a schema.
Note that all of these methods **must** be defined on a non-pointer object for them to be called.
Take the following simplified example of a `CompactDate` that only includes the Year and Month:
```go
type CompactDate struct {
Year int
Month int
}
func (d *CompactDate) UnmarshalJSON(data []byte) error {
if len(data) != 9 {
return errors.New("invalid compact date length")
}
var err error
d.Year, err = strconv.Atoi(string(data[1:5]))
if err != nil {
return err
}
d.Month, err = strconv.Atoi(string(data[7:8]))
if err != nil {
return err
}
return nil
}
func (d *CompactDate) MarshalJSON() ([]byte, error) {
buf := new(bytes.Buffer)
buf.WriteByte('"')
buf.WriteString(fmt.Sprintf("%d-%02d", d.Year, d.Month))
buf.WriteByte('"')
return buf.Bytes(), nil
}
func (CompactDate) JSONSchema() *Schema {
return &Schema{
Type: "string",
Title: "Compact Date",
Description: "Short date that only includes year and month",
Pattern: "^[0-9]{4}-[0-1][0-9]$",
}
}
```
The resulting schema generated for this struct would look like:
```json
{
"$schema": "http://json-schema.org/draft/2020-12/schema",
"$ref": "#/$defs/CompactDate",
"$defs": {
"CompactDate": {
"pattern": "^[0-9]{4}-[0-1][0-9]$",
"type": "string",
"title": "Compact Date",
"description": "Short date that only includes year and month"
}
}
}
```
+76
View File
@@ -0,0 +1,76 @@
package jsonschema
import (
"errors"
"fmt"
"net/url"
"strings"
)
// ID represents a Schema ID type which should always be a URI.
// See draft-bhutton-json-schema-00 section 8.2.1
type ID string
// EmptyID is used to explicitly define an ID with no value.
const EmptyID ID = ""
// Validate is used to check if the ID looks like a proper schema.
// This is done by parsing the ID as a URL and checking it has all the
// relevant parts.
func (id ID) Validate() error {
u, err := url.Parse(id.String())
if err != nil {
return fmt.Errorf("invalid URL: %w", err)
}
if u.Hostname() == "" {
return errors.New("missing hostname")
}
if !strings.Contains(u.Hostname(), ".") {
return errors.New("hostname does not look valid")
}
if u.Path == "" {
return errors.New("path is expected")
}
if u.Scheme != "https" && u.Scheme != "http" {
return errors.New("unexpected schema")
}
return nil
}
// Anchor sets the anchor part of the schema URI.
func (id ID) Anchor(name string) ID {
b := id.Base()
return ID(b.String() + "#" + name)
}
// Def adds or replaces a definition identifier.
func (id ID) Def(name string) ID {
b := id.Base()
return ID(b.String() + "#/$defs/" + name)
}
// Add appends the provided path to the id, and removes any
// anchor data that might be there.
func (id ID) Add(path string) ID {
b := id.Base()
if !strings.HasPrefix(path, "/") {
path = "/" + path
}
return ID(b.String() + path)
}
// Base removes any anchor information from the schema
func (id ID) Base() ID {
s := id.String()
i := strings.LastIndex(s, "#")
if i != -1 {
s = s[0:i]
}
s = strings.TrimRight(s, "/")
return ID(s)
}
// String provides string version of ID
func (id ID) String() string {
return string(id)
}
+1148
View File
File diff suppressed because it is too large. Load diff
+146
View File
@@ -0,0 +1,146 @@
package jsonschema
import (
"fmt"
"io/fs"
gopath "path"
"path/filepath"
"reflect"
"strings"
"go/ast"
"go/doc"
"go/parser"
"go/token"
)
type commentOptions struct {
fullObjectText bool // use the first sentence only?
}
// CommentOption allows for special configuration options when preparing Go
// source files for comment extraction.
type CommentOption func(*commentOptions)
// WithFullComment will configure the comment extraction to process to use an
// object type's full comment text instead of just the synopsis.
func WithFullComment() CommentOption {
return func(o *commentOptions) {
o.fullObjectText = true
}
}
// AddGoComments will update the reflectors comment map with all the comments
// found in the provided source directories including sub-directories, in order to
// generate a dictionary of comments associated with Types and Fields. The results
// will be added to the `Reflect.CommentMap` ready to use with Schema "description"
// fields.
//
// The `go/parser` library is used to extract all the comments and unfortunately doesn't
// have a built-in way to determine the fully qualified name of a package. The `base`
// parameter, the URL used to import that package, is thus required to be able to match
// reflected types.
//
// When parsing type comments, by default we use the `go/doc`'s Synopsis method to extract
// the first phrase only. Field comments, which tend to be much shorter, will include everything.
// This behavior can be changed by using the `WithFullComment` option.
func (r *Reflector) AddGoComments(base, path string, opts ...CommentOption) error {
if r.CommentMap == nil {
r.CommentMap = make(map[string]string)
}
co := new(commentOptions)
for _, opt := range opts {
opt(co)
}
return r.extractGoComments(base, path, r.CommentMap, co)
}
func (r *Reflector) extractGoComments(base, path string, commentMap map[string]string, opts *commentOptions) error {
fset := token.NewFileSet()
dict := make(map[string][]*ast.Package)
err := filepath.Walk(path, func(path string, info fs.FileInfo, err error) error {
if err != nil {
return err
}
if info.IsDir() {
d, err := parser.ParseDir(fset, path, nil, parser.ParseComments)
if err != nil {
return err
}
for _, v := range d {
// paths may have multiple packages, like for tests
k := gopath.Join(base, path)
dict[k] = append(dict[k], v)
}
}
return nil
})
if err != nil {
return err
}
for pkg, p := range dict {
for _, f := range p {
gtxt := ""
typ := ""
ast.Inspect(f, func(n ast.Node) bool {
switch x := n.(type) {
case *ast.TypeSpec:
typ = x.Name.String()
if !ast.IsExported(typ) {
typ = ""
} else {
txt := x.Doc.Text()
if txt == "" && gtxt != "" {
txt = gtxt
gtxt = ""
}
if !opts.fullObjectText {
txt = doc.Synopsis(txt)
}
commentMap[fmt.Sprintf("%s.%s", pkg, typ)] = strings.TrimSpace(txt)
}
case *ast.Field:
txt := x.Doc.Text()
if txt == "" {
txt = x.Comment.Text()
}
if typ != "" && txt != "" {
for _, n := range x.Names {
if ast.IsExported(n.String()) {
k := fmt.Sprintf("%s.%s.%s", pkg, typ, n)
commentMap[k] = strings.TrimSpace(txt)
}
}
}
case *ast.GenDecl:
// remember for the next type
gtxt = x.Doc.Text()
}
return true
})
}
}
return nil
}
func (r *Reflector) lookupComment(t reflect.Type, name string) string {
if r.LookupComment != nil {
if comment := r.LookupComment(t, name); comment != "" {
return comment
}
}
if r.CommentMap == nil {
return ""
}
n := fullyQualifiedTypeName(t)
if name != "" {
n = n + "." + name
}
return r.CommentMap[n]
}
+94
View File
@@ -0,0 +1,94 @@
package jsonschema
import (
"encoding/json"
orderedmap "github.com/wk8/go-ordered-map/v2"
)
// Version is the JSON Schema version.
var Version = "https://json-schema.org/draft/2020-12/schema"
// Schema represents a JSON Schema object type.
// RFC draft-bhutton-json-schema-00 section 4.3
type Schema struct {
// RFC draft-bhutton-json-schema-00
Version string `json:"$schema,omitempty"` // section 8.1.1
ID ID `json:"$id,omitempty"` // section 8.2.1
Anchor string `json:"$anchor,omitempty"` // section 8.2.2
Ref string `json:"$ref,omitempty"` // section 8.2.3.1
DynamicRef string `json:"$dynamicRef,omitempty"` // section 8.2.3.2
Definitions Definitions `json:"$defs,omitempty"` // section 8.2.4
Comments string `json:"$comment,omitempty"` // section 8.3
// RFC draft-bhutton-json-schema-00 section 10.2.1 (Sub-schemas with logic)
AllOf []*Schema `json:"allOf,omitempty"` // section 10.2.1.1
AnyOf []*Schema `json:"anyOf,omitempty"` // section 10.2.1.2
OneOf []*Schema `json:"oneOf,omitempty"` // section 10.2.1.3
Not *Schema `json:"not,omitempty"` // section 10.2.1.4
// RFC draft-bhutton-json-schema-00 section 10.2.2 (Apply sub-schemas conditionally)
If *Schema `json:"if,omitempty"` // section 10.2.2.1
Then *Schema `json:"then,omitempty"` // section 10.2.2.2
Else *Schema `json:"else,omitempty"` // section 10.2.2.3
DependentSchemas map[string]*Schema `json:"dependentSchemas,omitempty"` // section 10.2.2.4
// RFC draft-bhutton-json-schema-00 section 10.3.1 (arrays)
PrefixItems []*Schema `json:"prefixItems,omitempty"` // section 10.3.1.1
Items *Schema `json:"items,omitempty"` // section 10.3.1.2 (replaces additionalItems)
Contains *Schema `json:"contains,omitempty"` // section 10.3.1.3
// RFC draft-bhutton-json-schema-00 section 10.3.2 (sub-schemas)
Properties *orderedmap.OrderedMap[string, *Schema] `json:"properties,omitempty"` // section 10.3.2.1
PatternProperties map[string]*Schema `json:"patternProperties,omitempty"` // section 10.3.2.2
AdditionalProperties *Schema `json:"additionalProperties,omitempty"` // section 10.3.2.3
PropertyNames *Schema `json:"propertyNames,omitempty"` // section 10.3.2.4
// RFC draft-bhutton-json-schema-validation-00, section 6
Type string `json:"type,omitempty"` // section 6.1.1
Enum []any `json:"enum,omitempty"` // section 6.1.2
Const any `json:"const,omitempty"` // section 6.1.3
MultipleOf json.Number `json:"multipleOf,omitempty"` // section 6.2.1
Maximum json.Number `json:"maximum,omitempty"` // section 6.2.2
ExclusiveMaximum json.Number `json:"exclusiveMaximum,omitempty"` // section 6.2.3
Minimum json.Number `json:"minimum,omitempty"` // section 6.2.4
ExclusiveMinimum json.Number `json:"exclusiveMinimum,omitempty"` // section 6.2.5
MaxLength *uint64 `json:"maxLength,omitempty"` // section 6.3.1
MinLength *uint64 `json:"minLength,omitempty"` // section 6.3.2
Pattern string `json:"pattern,omitempty"` // section 6.3.3
MaxItems *uint64 `json:"maxItems,omitempty"` // section 6.4.1
MinItems *uint64 `json:"minItems,omitempty"` // section 6.4.2
UniqueItems bool `json:"uniqueItems,omitempty"` // section 6.4.3
MaxContains *uint64 `json:"maxContains,omitempty"` // section 6.4.4
MinContains *uint64 `json:"minContains,omitempty"` // section 6.4.5
MaxProperties *uint64 `json:"maxProperties,omitempty"` // section 6.5.1
MinProperties *uint64 `json:"minProperties,omitempty"` // section 6.5.2
Required []string `json:"required,omitempty"` // section 6.5.3
DependentRequired map[string][]string `json:"dependentRequired,omitempty"` // section 6.5.4
// RFC draft-bhutton-json-schema-validation-00, section 7
Format string `json:"format,omitempty"`
// RFC draft-bhutton-json-schema-validation-00, section 8
ContentEncoding string `json:"contentEncoding,omitempty"` // section 8.3
ContentMediaType string `json:"contentMediaType,omitempty"` // section 8.4
ContentSchema *Schema `json:"contentSchema,omitempty"` // section 8.5
// RFC draft-bhutton-json-schema-validation-00, section 9
Title string `json:"title,omitempty"` // section 9.1
Description string `json:"description,omitempty"` // section 9.1
Default any `json:"default,omitempty"` // section 9.2
Deprecated bool `json:"deprecated,omitempty"` // section 9.3
ReadOnly bool `json:"readOnly,omitempty"` // section 9.4
WriteOnly bool `json:"writeOnly,omitempty"` // section 9.4
Examples []any `json:"examples,omitempty"` // section 9.5
Extras map[string]any `json:"-"`
// Special boolean representation of the Schema - section 4.3.2
boolean *bool
}
var (
// TrueSchema defines a schema with a true value
TrueSchema = &Schema{boolean: &[]bool{true}[0]}
// FalseSchema defines a schema with a false value
FalseSchema = &Schema{boolean: &[]bool{false}[0]}
)
// Definitions hold schema definitions.
// http://json-schema.org/latest/json-schema-validation.html#rfc.section.5.26
// RFC draft-wright-json-schema-validation-00, section 5.26
type Definitions map[string]*Schema
+26
View File
@@ -0,0 +1,26 @@
package jsonschema
import (
"regexp"
"strings"
orderedmap "github.com/wk8/go-ordered-map/v2"
)
var matchFirstCap = regexp.MustCompile("(.)([A-Z][a-z]+)")
var matchAllCap = regexp.MustCompile("([a-z0-9])([A-Z])")
// ToSnakeCase converts the provided string into snake case using dashes.
// This is useful for Schema IDs and definitions to be coherent with
// common JSON Schema examples.
func ToSnakeCase(str string) string {
snake := matchFirstCap.ReplaceAllString(str, "${1}-${2}")
snake = matchAllCap.ReplaceAllString(snake, "${1}-${2}")
return strings.ToLower(snake)
}
// NewProperties is a helper method to instantiate a new properties ordered
// map.
func NewProperties() *orderedmap.OrderedMap[string, *Schema] {
return orderedmap.New[string, *Schema]()
}
+7
View File
@@ -0,0 +1,7 @@
Copyright (c) 2016 Mail.Ru Group
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+278
View File
@@ -0,0 +1,278 @@
// Package buffer implements a buffer for serialization, consisting of a chain of []byte-s to
// reduce copying and to allow reuse of individual chunks.
package buffer
import (
"io"
"net"
"sync"
)
// PoolConfig contains configuration for the allocation and reuse strategy.
type PoolConfig struct {
StartSize int // Minimum chunk size that is allocated.
PooledSize int // Minimum chunk size that is reused, reusing chunks too small will result in overhead.
MaxSize int // Maximum chunk size that will be allocated.
}
var config = PoolConfig{
StartSize: 128,
PooledSize: 512,
MaxSize: 32768,
}
// Reuse pool: chunk size -> pool.
var buffers = map[int]*sync.Pool{}
func initBuffers() {
for l := config.PooledSize; l <= config.MaxSize; l *= 2 {
buffers[l] = new(sync.Pool)
}
}
func init() {
initBuffers()
}
// Init sets up a non-default pooling and allocation strategy. Should be run before serialization is done.
func Init(cfg PoolConfig) {
config = cfg
initBuffers()
}
// putBuf puts a chunk to reuse pool if it can be reused.
func putBuf(buf []byte) {
size := cap(buf)
if size < config.PooledSize {
return
}
if c := buffers[size]; c != nil {
c.Put(buf[:0])
}
}
// getBuf gets a chunk from reuse pool or creates a new one if reuse failed.
func getBuf(size int) []byte {
if size >= config.PooledSize {
if c := buffers[size]; c != nil {
v := c.Get()
if v != nil {
return v.([]byte)
}
}
}
return make([]byte, 0, size)
}
// Buffer is a buffer optimized for serialization without extra copying.
type Buffer struct {
// Buf is the current chunk that can be used for serialization.
Buf []byte
toPool []byte
bufs [][]byte
}
// EnsureSpace makes sure that the current chunk contains at least s free bytes,
// possibly creating a new chunk.
func (b *Buffer) EnsureSpace(s int) {
if cap(b.Buf)-len(b.Buf) < s {
b.ensureSpaceSlow(s)
}
}
func (b *Buffer) ensureSpaceSlow(s int) {
l := len(b.Buf)
if l > 0 {
if cap(b.toPool) != cap(b.Buf) {
// Chunk was reallocated, toPool can be pooled.
putBuf(b.toPool)
}
if cap(b.bufs) == 0 {
b.bufs = make([][]byte, 0, 8)
}
b.bufs = append(b.bufs, b.Buf)
l = cap(b.toPool) * 2
} else {
l = config.StartSize
}
if l > config.MaxSize {
l = config.MaxSize
}
b.Buf = getBuf(l)
b.toPool = b.Buf
}
// AppendByte appends a single byte to buffer.
func (b *Buffer) AppendByte(data byte) {
b.EnsureSpace(1)
b.Buf = append(b.Buf, data)
}
// AppendBytes appends a byte slice to buffer.
func (b *Buffer) AppendBytes(data []byte) {
if len(data) <= cap(b.Buf)-len(b.Buf) {
b.Buf = append(b.Buf, data...) // fast path
} else {
b.appendBytesSlow(data)
}
}
func (b *Buffer) appendBytesSlow(data []byte) {
for len(data) > 0 {
b.EnsureSpace(1)
sz := cap(b.Buf) - len(b.Buf)
if sz > len(data) {
sz = len(data)
}
b.Buf = append(b.Buf, data[:sz]...)
data = data[sz:]
}
}
// AppendString appends a string to buffer.
func (b *Buffer) AppendString(data string) {
if len(data) <= cap(b.Buf)-len(b.Buf) {
b.Buf = append(b.Buf, data...) // fast path
} else {
b.appendStringSlow(data)
}
}
func (b *Buffer) appendStringSlow(data string) {
for len(data) > 0 {
b.EnsureSpace(1)
sz := cap(b.Buf) - len(b.Buf)
if sz > len(data) {
sz = len(data)
}
b.Buf = append(b.Buf, data[:sz]...)
data = data[sz:]
}
}
// Size computes the size of a buffer by adding sizes of every chunk.
func (b *Buffer) Size() int {
size := len(b.Buf)
for _, buf := range b.bufs {
size += len(buf)
}
return size
}
// DumpTo outputs the contents of a buffer to a writer and resets the buffer.
func (b *Buffer) DumpTo(w io.Writer) (written int, err error) {
bufs := net.Buffers(b.bufs)
if len(b.Buf) > 0 {
bufs = append(bufs, b.Buf)
}
n, err := bufs.WriteTo(w)
for _, buf := range b.bufs {
putBuf(buf)
}
putBuf(b.toPool)
b.bufs = nil
b.Buf = nil
b.toPool = nil
return int(n), err
}
// BuildBytes creates a single byte slice with all the contents of the buffer. Data is
// copied if it does not fit in a single chunk. You can optionally provide one byte
// slice as argument that it will try to reuse.
func (b *Buffer) BuildBytes(reuse ...[]byte) []byte {
if len(b.bufs) == 0 {
ret := b.Buf
b.toPool = nil
b.Buf = nil
return ret
}
var ret []byte
size := b.Size()
// If we got a buffer as argument and it is big enough, reuse it.
if len(reuse) == 1 && cap(reuse[0]) >= size {
ret = reuse[0][:0]
} else {
ret = make([]byte, 0, size)
}
for _, buf := range b.bufs {
ret = append(ret, buf...)
putBuf(buf)
}
ret = append(ret, b.Buf...)
putBuf(b.toPool)
b.bufs = nil
b.toPool = nil
b.Buf = nil
return ret
}
type readCloser struct {
offset int
bufs [][]byte
}
func (r *readCloser) Read(p []byte) (n int, err error) {
for _, buf := range r.bufs {
// Copy as much as we can.
x := copy(p[n:], buf[r.offset:])
n += x // Increment how much we filled.
// Did we empty the whole buffer?
if r.offset+x == len(buf) {
// On to the next buffer.
r.offset = 0
r.bufs = r.bufs[1:]
// We can release this buffer.
putBuf(buf)
} else {
r.offset += x
}
if n == len(p) {
break
}
}
// No buffers left or nothing read?
if len(r.bufs) == 0 {
err = io.EOF
}
return
}
func (r *readCloser) Close() error {
// Release all remaining buffers.
for _, buf := range r.bufs {
putBuf(buf)
}
// In case Close gets called multiple times.
r.bufs = nil
return nil
}
// ReadCloser creates an io.ReadCloser with all the contents of the buffer.
func (b *Buffer) ReadCloser() io.ReadCloser {
ret := &readCloser{0, append(b.bufs, b.Buf)}
b.bufs = nil
b.toPool = nil
b.Buf = nil
return ret
}
+405
View File
@@ -0,0 +1,405 @@
// Package jwriter contains a JSON writer.
package jwriter
import (
"io"
"strconv"
"unicode/utf8"
"github.com/mailru/easyjson/buffer"
)
// Flags describe various encoding options. The behavior may be actually implemented in the encoder, but
// Flags field in Writer is used to set and pass them around.
type Flags int
const (
NilMapAsEmpty Flags = 1 << iota // Encode nil map as '{}' rather than 'null'.
NilSliceAsEmpty // Encode nil slice as '[]' rather than 'null'.
)
// Writer is a JSON writer.
type Writer struct {
Flags Flags
Error error
Buffer buffer.Buffer
NoEscapeHTML bool
}
// Size returns the size of the data that was written out.
func (w *Writer) Size() int {
return w.Buffer.Size()
}
// DumpTo outputs the data to given io.Writer, resetting the buffer.
func (w *Writer) DumpTo(out io.Writer) (written int, err error) {
return w.Buffer.DumpTo(out)
}
// BuildBytes returns writer data as a single byte slice. You can optionally provide one byte slice
// as argument that it will try to reuse.
func (w *Writer) BuildBytes(reuse ...[]byte) ([]byte, error) {
if w.Error != nil {
return nil, w.Error
}
return w.Buffer.BuildBytes(reuse...), nil
}
// ReadCloser returns an io.ReadCloser that can be used to read the data.
// ReadCloser also resets the buffer.
func (w *Writer) ReadCloser() (io.ReadCloser, error) {
if w.Error != nil {
return nil, w.Error
}
return w.Buffer.ReadCloser(), nil
}
// RawByte appends raw binary data to the buffer.
func (w *Writer) RawByte(c byte) {
w.Buffer.AppendByte(c)
}
// RawByte appends raw binary data to the buffer.
func (w *Writer) RawString(s string) {
w.Buffer.AppendString(s)
}
// Raw appends raw binary data to the buffer or sets the error if it is given. Useful for
// calling with results of MarshalJSON-like functions.
func (w *Writer) Raw(data []byte, err error) {
switch {
case w.Error != nil:
return
case err != nil:
w.Error = err
case len(data) > 0:
w.Buffer.AppendBytes(data)
default:
w.RawString("null")
}
}
// RawText encloses raw binary data in quotes and appends in to the buffer.
// Useful for calling with results of MarshalText-like functions.
func (w *Writer) RawText(data []byte, err error) {
switch {
case w.Error != nil:
return
case err != nil:
w.Error = err
case len(data) > 0:
w.String(string(data))
default:
w.RawString("null")
}
}
// Base64Bytes appends data to the buffer after base64 encoding it
func (w *Writer) Base64Bytes(data []byte) {
if data == nil {
w.Buffer.AppendString("null")
return
}
w.Buffer.AppendByte('"')
w.base64(data)
w.Buffer.AppendByte('"')
}
func (w *Writer) Uint8(n uint8) {
w.Buffer.EnsureSpace(3)
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}
func (w *Writer) Uint16(n uint16) {
w.Buffer.EnsureSpace(5)
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}
func (w *Writer) Uint32(n uint32) {
w.Buffer.EnsureSpace(10)
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}
func (w *Writer) Uint(n uint) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}
func (w *Writer) Uint64(n uint64) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
}
func (w *Writer) Int8(n int8) {
w.Buffer.EnsureSpace(4)
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}
func (w *Writer) Int16(n int16) {
w.Buffer.EnsureSpace(6)
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}
func (w *Writer) Int32(n int32) {
w.Buffer.EnsureSpace(11)
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}
func (w *Writer) Int(n int) {
w.Buffer.EnsureSpace(21)
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}
func (w *Writer) Int64(n int64) {
w.Buffer.EnsureSpace(21)
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
}
func (w *Writer) Uint8Str(n uint8) {
w.Buffer.EnsureSpace(3)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Uint16Str(n uint16) {
w.Buffer.EnsureSpace(5)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Uint32Str(n uint32) {
w.Buffer.EnsureSpace(10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) UintStr(n uint) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Uint64Str(n uint64) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) UintptrStr(n uintptr) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Int8Str(n int8) {
w.Buffer.EnsureSpace(4)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Int16Str(n int16) {
w.Buffer.EnsureSpace(6)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Int32Str(n int32) {
w.Buffer.EnsureSpace(11)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) IntStr(n int) {
w.Buffer.EnsureSpace(21)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Int64Str(n int64) {
w.Buffer.EnsureSpace(21)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Float32(n float32) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
}
func (w *Writer) Float32Str(n float32) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Float64(n float64) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, n, 'g', -1, 64)
}
func (w *Writer) Float64Str(n float64) {
w.Buffer.EnsureSpace(20)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 64)
w.Buffer.Buf = append(w.Buffer.Buf, '"')
}
func (w *Writer) Bool(v bool) {
w.Buffer.EnsureSpace(5)
if v {
w.Buffer.Buf = append(w.Buffer.Buf, "true"...)
} else {
w.Buffer.Buf = append(w.Buffer.Buf, "false"...)
}
}
const chars = "0123456789abcdef"
func getTable(falseValues ...int) [128]bool {
table := [128]bool{}
for i := 0; i < 128; i++ {
table[i] = true
}
for _, v := range falseValues {
table[v] = false
}
return table
}
var (
htmlEscapeTable = getTable(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, '"', '&', '<', '>', '\\')
htmlNoEscapeTable = getTable(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, '"', '\\')
)
func (w *Writer) String(s string) {
w.Buffer.AppendByte('"')
// Portions of the string that contain no escapes are appended as
// byte slices.
p := 0 // last non-escape symbol
escapeTable := &htmlEscapeTable
if w.NoEscapeHTML {
escapeTable = &htmlNoEscapeTable
}
for i := 0; i < len(s); {
c := s[i]
if c < utf8.RuneSelf {
if escapeTable[c] {
// single-width character, no escaping is required
i++
continue
}
w.Buffer.AppendString(s[p:i])
switch c {
case '\t':
w.Buffer.AppendString(`\t`)
case '\r':
w.Buffer.AppendString(`\r`)
case '\n':
w.Buffer.AppendString(`\n`)
case '\\':
w.Buffer.AppendString(`\\`)
case '"':
w.Buffer.AppendString(`\"`)
default:
w.Buffer.AppendString(`\u00`)
w.Buffer.AppendByte(chars[c>>4])
w.Buffer.AppendByte(chars[c&0xf])
}
i++
p = i
continue
}
// broken utf
runeValue, runeWidth := utf8.DecodeRuneInString(s[i:])
if runeValue == utf8.RuneError && runeWidth == 1 {
w.Buffer.AppendString(s[p:i])
w.Buffer.AppendString(`\ufffd`)
i++
p = i
continue
}
// jsonp stuff - tab separator and line separator
if runeValue == '\u2028' || runeValue == '\u2029' {
w.Buffer.AppendString(s[p:i])
w.Buffer.AppendString(`\u202`)
w.Buffer.AppendByte(chars[runeValue&0xf])
i += runeWidth
p = i
continue
}
i += runeWidth
}
w.Buffer.AppendString(s[p:])
w.Buffer.AppendByte('"')
}
const encode = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
const padChar = '='
func (w *Writer) base64(in []byte) {
if len(in) == 0 {
return
}
w.Buffer.EnsureSpace(((len(in)-1)/3 + 1) * 4)
si := 0
n := (len(in) / 3) * 3
for si < n {
// Convert 3x 8bit source bytes into 4 bytes
val := uint(in[si+0])<<16 | uint(in[si+1])<<8 | uint(in[si+2])
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>18&0x3F], encode[val>>12&0x3F], encode[val>>6&0x3F], encode[val&0x3F])
si += 3
}
remain := len(in) - si
if remain == 0 {
return
}
// Add the remaining small block
val := uint(in[si+0]) << 16
if remain == 2 {
val |= uint(in[si+1]) << 8
}
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>18&0x3F], encode[val>>12&0x3F])
switch remain {
case 2:
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>6&0x3F], byte(padChar))
case 1:
w.Buffer.Buf = append(w.Buffer.Buf, byte(padChar), byte(padChar))
}
}
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2024 Anthropic, PBC
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+710
View File
@@ -0,0 +1,710 @@
package client
import (
"context"
"encoding/json"
"fmt"
"net/http"
"slices"
"sync"
"sync/atomic"
"github.com/mark3labs/mcp-go/client/transport"
"github.com/mark3labs/mcp-go/mcp"
)
// Client implements the MCP client.
type Client struct {
transport transport.Interface
initialized bool
notifications []func(mcp.JSONRPCNotification)
notifyMu sync.RWMutex
requestID atomic.Int64
clientCapabilities mcp.ClientCapabilities
serverCapabilities mcp.ServerCapabilities
protocolVersion string
samplingHandler SamplingHandler
rootsHandler RootsHandler
elicitationHandler ElicitationHandler
}
type ClientOption func(*Client)
// WithClientCapabilities sets the client capabilities for the client.
func WithClientCapabilities(capabilities mcp.ClientCapabilities) ClientOption {
return func(c *Client) {
c.clientCapabilities = capabilities
}
}
// WithSamplingHandler sets the sampling handler for the client.
// When set, the client will declare sampling capability during initialization.
func WithSamplingHandler(handler SamplingHandler) ClientOption {
return func(c *Client) {
c.samplingHandler = handler
}
}
// WithRootsHandler sets the roots handler for the client.
// WithRootsHandler returns a ClientOption that sets the client's RootsHandler.
// When provided, the client will declare the roots capability (ListChanged) during initialization.
func WithRootsHandler(handler RootsHandler) ClientOption {
return func(c *Client) {
c.rootsHandler = handler
}
}
// WithElicitationHandler sets the elicitation handler for the client.
// When set, the client will declare elicitation capability during initialization.
func WithElicitationHandler(handler ElicitationHandler) ClientOption {
return func(c *Client) {
c.elicitationHandler = handler
}
}
// WithSession assumes a MCP Session has already been initialized
func WithSession() ClientOption {
return func(c *Client) {
c.initialized = true
}
}
// NewClient creates a new MCP client with the given transport.
// Usage:
//
// stdio := transport.NewStdio("./mcp_server", nil, "xxx")
// client, err := NewClient(stdio)
// if err != nil {
// log.Fatalf("Failed to create client: %v", err)
// }
func NewClient(transport transport.Interface, options ...ClientOption) *Client {
client := &Client{
transport: transport,
}
for _, opt := range options {
opt(client)
}
return client
}
// Start initiates the connection to the server.
// Must be called before using the client.
func (c *Client) Start(ctx context.Context) error {
if c.transport == nil {
return fmt.Errorf("transport is nil")
}
// Start is idempotent - transports handle being called multiple times
err := c.transport.Start(ctx)
if err != nil {
return err
}
c.transport.SetNotificationHandler(func(notification mcp.JSONRPCNotification) {
c.notifyMu.RLock()
defer c.notifyMu.RUnlock()
for _, handler := range c.notifications {
handler(notification)
}
})
// Set up request handler for bidirectional communication (e.g., sampling)
if bidirectional, ok := c.transport.(transport.BidirectionalInterface); ok {
bidirectional.SetRequestHandler(c.handleIncomingRequest)
}
return nil
}
// Close shuts down the client and closes the transport.
func (c *Client) Close() error {
return c.transport.Close()
}
// OnNotification registers a handler function to be called when notifications are received.
// Multiple handlers can be registered and will be called in the order they were added.
func (c *Client) OnNotification(
handler func(notification mcp.JSONRPCNotification),
) {
c.notifyMu.Lock()
defer c.notifyMu.Unlock()
c.notifications = append(c.notifications, handler)
}
// OnConnectionLost registers a handler function to be called when the connection is lost.
// This is useful for handling HTTP2 idle timeout disconnections that should not be treated as errors.
func (c *Client) OnConnectionLost(handler func(error)) {
type connectionLostSetter interface {
SetConnectionLostHandler(func(error))
}
if setter, ok := c.transport.(connectionLostSetter); ok {
setter.SetConnectionLostHandler(handler)
}
}
// sendRequest sends a JSON-RPC request to the server and waits for a response.
// Returns the raw JSON response message or an error if the request fails.
func (c *Client) sendRequest(
ctx context.Context,
method string,
params any,
header http.Header,
) (*json.RawMessage, error) {
if !c.initialized && method != "initialize" {
return nil, fmt.Errorf("client not initialized")
}
id := c.requestID.Add(1)
request := transport.JSONRPCRequest{
JSONRPC: mcp.JSONRPC_VERSION,
ID: mcp.NewRequestId(id),
Method: method,
Params: params,
Header: header,
}
response, err := c.transport.SendRequest(ctx, request)
if err != nil {
return nil, transport.NewError(err)
}
if response.Error != nil {
return nil, response.Error.AsError()
}
return &response.Result, nil
}
// Initialize negotiates with the server.
// Must be called after Start, and before any request methods.
func (c *Client) Initialize(
ctx context.Context,
request mcp.InitializeRequest,
) (*mcp.InitializeResult, error) {
// Merge client capabilities with sampling capability if handler is configured
capabilities := request.Params.Capabilities
if c.samplingHandler != nil {
capabilities.Sampling = &struct{}{}
}
if c.rootsHandler != nil {
capabilities.Roots = &struct {
ListChanged bool `json:"listChanged,omitempty"`
}{
ListChanged: true,
}
}
// Add elicitation capability if handler is configured
if c.elicitationHandler != nil {
capabilities.Elicitation = &struct{}{}
}
// Ensure we send a params object with all required fields
params := struct {
ProtocolVersion string `json:"protocolVersion"`
ClientInfo mcp.Implementation `json:"clientInfo"`
Capabilities mcp.ClientCapabilities `json:"capabilities"`
}{
ProtocolVersion: request.Params.ProtocolVersion,
ClientInfo: request.Params.ClientInfo,
Capabilities: capabilities,
}
response, err := c.sendRequest(ctx, "initialize", params, request.Header)
if err != nil {
return nil, err
}
var result mcp.InitializeResult
if err := json.Unmarshal(*response, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
// Validate protocol version
if !slices.Contains(mcp.ValidProtocolVersions, result.ProtocolVersion) {
return nil, mcp.UnsupportedProtocolVersionError{Version: result.ProtocolVersion}
}
// Store serverCapabilities and protocol version
c.serverCapabilities = result.Capabilities
c.protocolVersion = result.ProtocolVersion
// Set protocol version on HTTP transports
if httpConn, ok := c.transport.(transport.HTTPConnection); ok {
httpConn.SetProtocolVersion(result.ProtocolVersion)
}
// Send initialized notification
notification := mcp.JSONRPCNotification{
JSONRPC: mcp.JSONRPC_VERSION,
Notification: mcp.Notification{
Method: "notifications/initialized",
},
}
err = c.transport.SendNotification(ctx, notification)
if err != nil {
return nil, fmt.Errorf(
"failed to send initialized notification: %w",
err,
)
}
c.initialized = true
return &result, nil
}
func (c *Client) Ping(ctx context.Context) error {
_, err := c.sendRequest(ctx, "ping", nil, nil)
return err
}
// ListResourcesByPage manually list resources by page.
func (c *Client) ListResourcesByPage(
ctx context.Context,
request mcp.ListResourcesRequest,
) (*mcp.ListResourcesResult, error) {
result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, "resources/list")
if err != nil {
return nil, err
}
return result, nil
}
func (c *Client) ListResources(
ctx context.Context,
request mcp.ListResourcesRequest,
) (*mcp.ListResourcesResult, error) {
result, err := c.ListResourcesByPage(ctx, request)
if err != nil {
return nil, err
}
for result.NextCursor != "" {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
request.Params.Cursor = result.NextCursor
newPageRes, err := c.ListResourcesByPage(ctx, request)
if err != nil {
return nil, err
}
result.Resources = append(result.Resources, newPageRes.Resources...)
result.NextCursor = newPageRes.NextCursor
}
}
return result, nil
}
func (c *Client) ListResourceTemplatesByPage(
ctx context.Context,
request mcp.ListResourceTemplatesRequest,
) (*mcp.ListResourceTemplatesResult, error) {
result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, "resources/templates/list")
if err != nil {
return nil, err
}
return result, nil
}
func (c *Client) ListResourceTemplates(
ctx context.Context,
request mcp.ListResourceTemplatesRequest,
) (*mcp.ListResourceTemplatesResult, error) {
result, err := c.ListResourceTemplatesByPage(ctx, request)
if err != nil {
return nil, err
}
for result.NextCursor != "" {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
request.Params.Cursor = result.NextCursor
newPageRes, err := c.ListResourceTemplatesByPage(ctx, request)
if err != nil {
return nil, err
}
result.ResourceTemplates = append(result.ResourceTemplates, newPageRes.ResourceTemplates...)
result.NextCursor = newPageRes.NextCursor
}
}
return result, nil
}
func (c *Client) ReadResource(
ctx context.Context,
request mcp.ReadResourceRequest,
) (*mcp.ReadResourceResult, error) {
response, err := c.sendRequest(ctx, "resources/read", request.Params, request.Header)
if err != nil {
return nil, err
}
return mcp.ParseReadResourceResult(response)
}
func (c *Client) Subscribe(
ctx context.Context,
request mcp.SubscribeRequest,
) error {
_, err := c.sendRequest(ctx, "resources/subscribe", request.Params, request.Header)
return err
}
func (c *Client) Unsubscribe(
ctx context.Context,
request mcp.UnsubscribeRequest,
) error {
_, err := c.sendRequest(ctx, "resources/unsubscribe", request.Params, request.Header)
return err
}
func (c *Client) ListPromptsByPage(
ctx context.Context,
request mcp.ListPromptsRequest,
) (*mcp.ListPromptsResult, error) {
result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, "prompts/list")
if err != nil {
return nil, err
}
return result, nil
}
func (c *Client) ListPrompts(
ctx context.Context,
request mcp.ListPromptsRequest,
) (*mcp.ListPromptsResult, error) {
result, err := c.ListPromptsByPage(ctx, request)
if err != nil {
return nil, err
}
for result.NextCursor != "" {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
request.Params.Cursor = result.NextCursor
newPageRes, err := c.ListPromptsByPage(ctx, request)
if err != nil {
return nil, err
}
result.Prompts = append(result.Prompts, newPageRes.Prompts...)
result.NextCursor = newPageRes.NextCursor
}
}
return result, nil
}
func (c *Client) GetPrompt(
ctx context.Context,
request mcp.GetPromptRequest,
) (*mcp.GetPromptResult, error) {
response, err := c.sendRequest(ctx, "prompts/get", request.Params, request.Header)
if err != nil {
return nil, err
}
return mcp.ParseGetPromptResult(response)
}
func (c *Client) ListToolsByPage(
ctx context.Context,
request mcp.ListToolsRequest,
) (*mcp.ListToolsResult, error) {
result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, "tools/list")
if err != nil {
return nil, err
}
return result, nil
}
func (c *Client) ListTools(
ctx context.Context,
request mcp.ListToolsRequest,
) (*mcp.ListToolsResult, error) {
result, err := c.ListToolsByPage(ctx, request)
if err != nil {
return nil, err
}
for result.NextCursor != "" {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
request.Params.Cursor = result.NextCursor
newPageRes, err := c.ListToolsByPage(ctx, request)
if err != nil {
return nil, err
}
result.Tools = append(result.Tools, newPageRes.Tools...)
result.NextCursor = newPageRes.NextCursor
}
}
return result, nil
}
func (c *Client) CallTool(
ctx context.Context,
request mcp.CallToolRequest,
) (*mcp.CallToolResult, error) {
response, err := c.sendRequest(ctx, "tools/call", request.Params, request.Header)
if err != nil {
return nil, err
}
return mcp.ParseCallToolResult(response)
}
func (c *Client) SetLevel(
ctx context.Context,
request mcp.SetLevelRequest,
) error {
_, err := c.sendRequest(ctx, "logging/setLevel", request.Params, request.Header)
return err
}
func (c *Client) Complete(
ctx context.Context,
request mcp.CompleteRequest,
) (*mcp.CompleteResult, error) {
response, err := c.sendRequest(ctx, "completion/complete", request.Params, request.Header)
if err != nil {
return nil, err
}
var result mcp.CompleteResult
if err := json.Unmarshal(*response, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
return &result, nil
}
// RootListChanges sends a roots list-changed notification to the server.
func (c *Client) RootListChanges(
ctx context.Context,
) error {
// Send root list changes notification
notification := mcp.JSONRPCNotification{
JSONRPC: mcp.JSONRPC_VERSION,
Notification: mcp.Notification{
Method: mcp.MethodNotificationRootsListChanged,
},
}
err := c.transport.SendNotification(ctx, notification)
if err != nil {
return fmt.Errorf(
"failed to send root list change notification: %w",
err,
)
}
return nil
}
// handleIncomingRequest processes incoming requests from the server.
// This is the main entry point for server-to-client requests like sampling and elicitation.
func (c *Client) handleIncomingRequest(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
switch request.Method {
case string(mcp.MethodSamplingCreateMessage):
return c.handleSamplingRequestTransport(ctx, request)
case string(mcp.MethodElicitationCreate):
return c.handleElicitationRequestTransport(ctx, request)
case string(mcp.MethodPing):
return c.handlePingRequestTransport(ctx, request)
case string(mcp.MethodListRoots):
return c.handleListRootsRequestTransport(ctx, request)
default:
return nil, fmt.Errorf("unsupported request method: %s", request.Method)
}
}
// handleSamplingRequestTransport handles sampling requests at the transport level.
func (c *Client) handleSamplingRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
if c.samplingHandler == nil {
return nil, fmt.Errorf("no sampling handler configured")
}
// Parse the request parameters
var params mcp.CreateMessageParams
if request.Params != nil {
paramsBytes, err := json.Marshal(request.Params)
if err != nil {
return nil, fmt.Errorf("failed to marshal params: %w", err)
}
if err := json.Unmarshal(paramsBytes, &params); err != nil {
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
}
}
// Fix content parsing - HTTP transport unmarshals TextContent as map[string]any
// Use the helper function to properly handle content from different transports
for i := range params.Messages {
if contentMap, ok := params.Messages[i].Content.(map[string]any); ok {
// Parse the content map into a proper Content type
content, err := mcp.ParseContent(contentMap)
if err != nil {
return nil, fmt.Errorf("failed to parse content for message %d: %w", i, err)
}
params.Messages[i].Content = content
}
}
// Create the MCP request
mcpRequest := mcp.CreateMessageRequest{
Request: mcp.Request{
Method: string(mcp.MethodSamplingCreateMessage),
},
CreateMessageParams: params,
}
// Call the sampling handler
result, err := c.samplingHandler.CreateMessage(ctx, mcpRequest)
if err != nil {
return nil, err
}
// Marshal the result
resultBytes, err := json.Marshal(result)
if err != nil {
return nil, fmt.Errorf("failed to marshal result: %w", err)
}
// Create the transport response
response := transport.NewJSONRPCResultResponse(request.ID, json.RawMessage(resultBytes))
return response, nil
}
// handleListRootsRequestTransport handles list roots requests at the transport level.
func (c *Client) handleListRootsRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
if c.rootsHandler == nil {
return nil, fmt.Errorf("no roots handler configured")
}
// Create the MCP request
mcpRequest := mcp.ListRootsRequest{
Request: mcp.Request{
Method: string(mcp.MethodListRoots),
},
}
// Call the list roots handler
result, err := c.rootsHandler.ListRoots(ctx, mcpRequest)
if err != nil {
return nil, err
}
// Marshal the result
resultBytes, err := json.Marshal(result)
if err != nil {
return nil, fmt.Errorf("failed to marshal result: %w", err)
}
// Create the transport response
response := transport.NewJSONRPCResultResponse(request.ID, json.RawMessage(resultBytes))
return response, nil
}
// handleElicitationRequestTransport handles elicitation requests at the transport level.
func (c *Client) handleElicitationRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
if c.elicitationHandler == nil {
return nil, fmt.Errorf("no elicitation handler configured")
}
// Parse the request parameters
var params mcp.ElicitationParams
if request.Params != nil {
paramsBytes, err := json.Marshal(request.Params)
if err != nil {
return nil, fmt.Errorf("failed to marshal params: %w", err)
}
if err := json.Unmarshal(paramsBytes, &params); err != nil {
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
}
}
// Create the MCP request
mcpRequest := mcp.ElicitationRequest{
Request: mcp.Request{
Method: string(mcp.MethodElicitationCreate),
},
Params: params,
}
// Call the elicitation handler
result, err := c.elicitationHandler.Elicit(ctx, mcpRequest)
if err != nil {
return nil, err
}
// Marshal the result
resultBytes, err := json.Marshal(result)
if err != nil {
return nil, fmt.Errorf("failed to marshal result: %w", err)
}
// Create the transport response
response := transport.NewJSONRPCResultResponse(request.ID, resultBytes)
return response, nil
}
func (c *Client) handlePingRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
b, _ := json.Marshal(&mcp.EmptyResult{})
return transport.NewJSONRPCResultResponse(request.ID, b), nil
}
func listByPage[T any](
ctx context.Context,
client *Client,
request mcp.PaginatedRequest,
method string,
) (*T, error) {
response, err := client.sendRequest(ctx, method, request.Params, nil)
if err != nil {
return nil, err
}
var result T
if err := json.Unmarshal(*response, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
return &result, nil
}
// Helper methods
// GetTransport gives access to the underlying transport layer.
// Cast it to the specific transport type and obtain the other helper methods.
func (c *Client) GetTransport() transport.Interface {
return c.transport
}
// GetServerCapabilities returns the server capabilities.
func (c *Client) GetServerCapabilities() mcp.ServerCapabilities {
return c.serverCapabilities
}
// GetClientCapabilities returns the client capabilities.
func (c *Client) GetClientCapabilities() mcp.ClientCapabilities {
return c.clientCapabilities
}
// GetSessionId returns the session ID of the transport.
// If the transport does not support sessions, it returns an empty string.
func (c *Client) GetSessionId() string {
if c.transport == nil {
return ""
}
return c.transport.GetSessionId()
}
// IsInitialized returns true if the client has been initialized.
func (c *Client) IsInitialized() bool {
return c.initialized
}
+19
View File
@@ -0,0 +1,19 @@
package client
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
)
// ElicitationHandler defines the interface for handling elicitation requests from servers.
// Clients can implement this interface to request additional information from users.
type ElicitationHandler interface {
// Elicit handles an elicitation request from the server and returns the user's response.
// The implementation should:
// 1. Present the request message to the user
// 2. Validate input against the requested schema
// 3. Allow the user to accept, decline, or cancel
// 4. Return the appropriate response
Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
}
+22
View File
@@ -0,0 +1,22 @@
package client
import (
"fmt"
"github.com/mark3labs/mcp-go/client/transport"
)
// NewStreamableHttpClient is a convenience method that creates a new streamable-http-based MCP client
// with the given base URL. Returns an error if the URL is invalid.
func NewStreamableHttpClient(baseURL string, options ...transport.StreamableHTTPCOption) (*Client, error) {
trans, err := transport.NewStreamableHTTP(baseURL, options...)
if err != nil {
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
}
clientOptions := make([]ClientOption, 0)
sessionID := trans.GetSessionId()
if sessionID != "" {
clientOptions = append(clientOptions, WithSession())
}
return NewClient(trans, clientOptions...), nil
}
+38
View File
@@ -0,0 +1,38 @@
package client
import (
"context"
"github.com/mark3labs/mcp-go/client/transport"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// NewInProcessClient connect directly to a mcp server object in the same process
func NewInProcessClient(server *server.MCPServer) (*Client, error) {
inProcessTransport := transport.NewInProcessTransport(server)
return NewClient(inProcessTransport), nil
}
// NewInProcessClientWithSamplingHandler creates an in-process client with sampling support
func NewInProcessClientWithSamplingHandler(server *server.MCPServer, handler SamplingHandler) (*Client, error) {
// Create a wrapper that implements server.SamplingHandler
serverHandler := &inProcessSamplingHandlerWrapper{handler: handler}
inProcessTransport := transport.NewInProcessTransportWithOptions(server,
transport.WithSamplingHandler(serverHandler))
client := NewClient(inProcessTransport)
client.samplingHandler = handler
return client, nil
}
// inProcessSamplingHandlerWrapper wraps client.SamplingHandler to implement server.SamplingHandler
type inProcessSamplingHandlerWrapper struct {
handler SamplingHandler
}
func (w *inProcessSamplingHandlerWrapper) CreateMessage(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
return w.handler.CreateMessage(ctx, request)
}
+109
View File
@@ -0,0 +1,109 @@
// Package client provides MCP (Model Context Protocol) client implementations.
package client
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
)
// MCPClient represents an MCP client interface
type MCPClient interface {
// Initialize sends the initial connection request to the server
Initialize(
ctx context.Context,
request mcp.InitializeRequest,
) (*mcp.InitializeResult, error)
// Ping checks if the server is alive
Ping(ctx context.Context) error
// ListResourcesByPage manually list resources by page.
ListResourcesByPage(
ctx context.Context,
request mcp.ListResourcesRequest,
) (*mcp.ListResourcesResult, error)
// ListResources requests a list of available resources from the server
ListResources(
ctx context.Context,
request mcp.ListResourcesRequest,
) (*mcp.ListResourcesResult, error)
// ListResourceTemplatesByPage manually list resource templates by page.
ListResourceTemplatesByPage(
ctx context.Context,
request mcp.ListResourceTemplatesRequest,
) (*mcp.ListResourceTemplatesResult,
error)
// ListResourceTemplates requests a list of available resource templates from the server
ListResourceTemplates(
ctx context.Context,
request mcp.ListResourceTemplatesRequest,
) (*mcp.ListResourceTemplatesResult,
error)
// ReadResource reads a specific resource from the server
ReadResource(
ctx context.Context,
request mcp.ReadResourceRequest,
) (*mcp.ReadResourceResult, error)
// Subscribe requests notifications for changes to a specific resource
Subscribe(ctx context.Context, request mcp.SubscribeRequest) error
// Unsubscribe cancels notifications for a specific resource
Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error
// ListPromptsByPage manually list prompts by page.
ListPromptsByPage(
ctx context.Context,
request mcp.ListPromptsRequest,
) (*mcp.ListPromptsResult, error)
// ListPrompts requests a list of available prompts from the server
ListPrompts(
ctx context.Context,
request mcp.ListPromptsRequest,
) (*mcp.ListPromptsResult, error)
// GetPrompt retrieves a specific prompt from the server
GetPrompt(
ctx context.Context,
request mcp.GetPromptRequest,
) (*mcp.GetPromptResult, error)
// ListToolsByPage manually list tools by page.
ListToolsByPage(
ctx context.Context,
request mcp.ListToolsRequest,
) (*mcp.ListToolsResult, error)
// ListTools requests a list of available tools from the server
ListTools(
ctx context.Context,
request mcp.ListToolsRequest,
) (*mcp.ListToolsResult, error)
// CallTool invokes a specific tool on the server
CallTool(
ctx context.Context,
request mcp.CallToolRequest,
) (*mcp.CallToolResult, error)
// SetLevel sets the logging level for the server
SetLevel(ctx context.Context, request mcp.SetLevelRequest) error
// Complete requests completion options for a given argument
Complete(
ctx context.Context,
request mcp.CompleteRequest,
) (*mcp.CompleteResult, error)
// Close client connection and cleanup resources
Close() error
// OnNotification registers a handler for notifications
OnNotification(handler func(notification mcp.JSONRPCNotification))
}
+76
View File
@@ -0,0 +1,76 @@
package client
import (
"errors"
"fmt"
"github.com/mark3labs/mcp-go/client/transport"
)
// OAuthConfig is a convenience type that wraps transport.OAuthConfig
type OAuthConfig = transport.OAuthConfig
// Token is a convenience type that wraps transport.Token
type Token = transport.Token
// TokenStore is a convenience type that wraps transport.TokenStore
type TokenStore = transport.TokenStore
// MemoryTokenStore is a convenience type that wraps transport.MemoryTokenStore
type MemoryTokenStore = transport.MemoryTokenStore
// NewMemoryTokenStore is a convenience function that wraps transport.NewMemoryTokenStore
var NewMemoryTokenStore = transport.NewMemoryTokenStore
// NewOAuthStreamableHttpClient creates a new streamable-http-based MCP client with OAuth support.
// Returns an error if the URL is invalid.
func NewOAuthStreamableHttpClient(baseURL string, oauthConfig OAuthConfig, options ...transport.StreamableHTTPCOption) (*Client, error) {
// Add OAuth option to the list of options
options = append(options, transport.WithHTTPOAuth(oauthConfig))
trans, err := transport.NewStreamableHTTP(baseURL, options...)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP transport: %w", err)
}
return NewClient(trans), nil
}
// NewOAuthStreamableHttpClient creates a new streamable-http-based MCP client with OAuth support.
// Returns an error if the URL is invalid.
func NewOAuthSSEClient(baseURL string, oauthConfig OAuthConfig, options ...transport.ClientOption) (*Client, error) {
// Add OAuth option to the list of options
options = append(options, transport.WithOAuth(oauthConfig))
trans, err := transport.NewSSE(baseURL, options...)
if err != nil {
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
}
return NewClient(trans), nil
}
// GenerateCodeVerifier generates a code verifier for PKCE
var GenerateCodeVerifier = transport.GenerateCodeVerifier
// GenerateCodeChallenge generates a code challenge from a code verifier
var GenerateCodeChallenge = transport.GenerateCodeChallenge
// GenerateState generates a state parameter for OAuth
var GenerateState = transport.GenerateState
// OAuthAuthorizationRequiredError is returned when OAuth authorization is required
type OAuthAuthorizationRequiredError = transport.OAuthAuthorizationRequiredError
// IsOAuthAuthorizationRequiredError checks if an error is an OAuthAuthorizationRequiredError
func IsOAuthAuthorizationRequiredError(err error) bool {
var target *OAuthAuthorizationRequiredError
return errors.As(err, &target)
}
// GetOAuthHandler extracts the OAuthHandler from an OAuthAuthorizationRequiredError
func GetOAuthHandler(err error) *transport.OAuthHandler {
var oauthErr *OAuthAuthorizationRequiredError
if errors.As(err, &oauthErr) {
return oauthErr.Handler
}
return nil
}
+17
View File
@@ -0,0 +1,17 @@
package client
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
)
// RootsHandler defines the interface for handling roots requests from servers.
// Clients can implement this interface to provide roots list to servers.
type RootsHandler interface {
// ListRoots handles a list root request from the server and returns the roots list.
// The implementation should:
// 1. Validate input against the requested schema
// 2. Return the appropriate response
ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error)
}
+20
View File
@@ -0,0 +1,20 @@
package client
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
)
// SamplingHandler defines the interface for handling sampling requests from servers.
// Clients can implement this interface to provide LLM sampling capabilities to servers.
type SamplingHandler interface {
// CreateMessage handles a sampling request from the server and returns the generated message.
// The implementation should:
// 1. Validate the request parameters
// 2. Optionally prompt the user for approval (human-in-the-loop)
// 3. Select an appropriate model based on preferences
// 4. Generate the response using the selected model
// 5. Return the result with model information and stop reason
CreateMessage(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error)
}
+40
View File
@@ -0,0 +1,40 @@
package client
import (
"fmt"
"net/http"
"net/url"
"github.com/mark3labs/mcp-go/client/transport"
)
func WithHeaders(headers map[string]string) transport.ClientOption {
return transport.WithHeaders(headers)
}
func WithHeaderFunc(headerFunc transport.HTTPHeaderFunc) transport.ClientOption {
return transport.WithHeaderFunc(headerFunc)
}
func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
return transport.WithHTTPClient(httpClient)
}
// NewSSEMCPClient creates a new SSE-based MCP client with the given base URL.
// Returns an error if the URL is invalid.
func NewSSEMCPClient(baseURL string, options ...transport.ClientOption) (*Client, error) {
sseTransport, err := transport.NewSSE(baseURL, options...)
if err != nil {
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
}
return NewClient(sseTransport), nil
}
// GetEndpoint returns the current endpoint URL for the SSE connection.
//
// Note: This method only works with SSE transport, or it will panic.
func GetEndpoint(c *Client) *url.URL {
t := c.GetTransport()
sse := t.(*transport.SSE)
return sse.GetEndpoint()
}
+58
View File
@@ -0,0 +1,58 @@
package client
import (
"context"
"fmt"
"io"
"github.com/mark3labs/mcp-go/client/transport"
)
// NewStdioMCPClient creates a new stdio-based MCP client that communicates with a subprocess.
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
// Returns an error if the subprocess cannot be started or the pipes cannot be created.
//
// NOTICE: NewStdioMCPClient will start the connection automatically.
// This is for backward compatibility.
func NewStdioMCPClient(
command string,
env []string,
args ...string,
) (*Client, error) {
return NewStdioMCPClientWithOptions(command, env, args)
}
// NewStdioMCPClientWithOptions creates a new stdio-based MCP client that communicates with a subprocess.
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
// Optional configuration functions can be provided to customize the transport before it starts,
// such as setting a custom command function.
//
// NOTICE: NewStdioMCPClientWithOptions automatically starts the underlying transport.
// This is for backward compatibility.
func NewStdioMCPClientWithOptions(
command string,
env []string,
args []string,
opts ...transport.StdioOption,
) (*Client, error) {
stdioTransport := transport.NewStdioWithOptions(command, env, args, opts...)
if err := stdioTransport.Start(context.Background()); err != nil {
return nil, fmt.Errorf("failed to start stdio transport: %w", err)
}
return NewClient(stdioTransport), nil
}
// GetStderr returns a reader for the stderr output of the subprocess.
// This can be used to capture error messages or logs from the subprocess.
func GetStderr(c *Client) (io.Reader, bool) {
t := c.GetTransport()
stdio, ok := t.(*transport.Stdio)
if !ok {
return nil, false
}
return stdio.Stderr(), true
}
+7
View File
@@ -0,0 +1,7 @@
package transport
// Common HTTP header constants used across transports
const (
HeaderKeySessionID = "Mcp-Session-Id"
HeaderKeyProtocolVersion = "Mcp-Protocol-Version"
)
+22
View File
@@ -0,0 +1,22 @@
package transport
import "fmt"
// Error wraps a low-level transport error in a concrete type.
type Error struct {
Err error
}
func (e *Error) Error() string {
return fmt.Sprintf("transport error: %v", e.Err)
}
func (e *Error) Unwrap() error {
return e.Err
}
func NewError(err error) *Error {
return &Error{
Err: err,
}
}
+140
View File
@@ -0,0 +1,140 @@
package transport
import (
"context"
"encoding/json"
"fmt"
"sync"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
type InProcessTransport struct {
server *server.MCPServer
samplingHandler server.SamplingHandler
elicitationHandler server.ElicitationHandler
rootsHandler server.RootsHandler
session *server.InProcessSession
sessionID string
onNotification func(mcp.JSONRPCNotification)
notifyMu sync.RWMutex
started bool
startedMu sync.Mutex
}
type InProcessOption func(*InProcessTransport)
func WithSamplingHandler(handler server.SamplingHandler) InProcessOption {
return func(t *InProcessTransport) {
t.samplingHandler = handler
}
}
func WithElicitationHandler(handler server.ElicitationHandler) InProcessOption {
return func(t *InProcessTransport) {
t.elicitationHandler = handler
}
}
func WithRootsHandler(handler server.RootsHandler) InProcessOption {
return func(t *InProcessTransport) {
t.rootsHandler = handler
}
}
func NewInProcessTransport(server *server.MCPServer) *InProcessTransport {
return &InProcessTransport{
server: server,
}
}
func NewInProcessTransportWithOptions(server *server.MCPServer, opts ...InProcessOption) *InProcessTransport {
t := &InProcessTransport{
server: server,
sessionID: server.GenerateInProcessSessionID(),
}
for _, opt := range opts {
opt(t)
}
return t
}
func (c *InProcessTransport) Start(ctx context.Context) error {
c.startedMu.Lock()
if c.started {
c.startedMu.Unlock()
return nil
}
c.started = true
c.startedMu.Unlock()
// Create and register session if we have handlers
if c.samplingHandler != nil || c.elicitationHandler != nil || c.rootsHandler != nil {
c.session = server.NewInProcessSessionWithHandlers(c.sessionID, c.samplingHandler, c.elicitationHandler, c.rootsHandler)
if err := c.server.RegisterSession(ctx, c.session); err != nil {
c.startedMu.Lock()
c.started = false
c.startedMu.Unlock()
return fmt.Errorf("failed to register session: %w", err)
}
}
return nil
}
func (c *InProcessTransport) SendRequest(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error) {
requestBytes, err := json.Marshal(request)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
requestBytes = append(requestBytes, '\n')
// Add session to context if available
if c.session != nil {
ctx = c.server.WithContext(ctx, c.session)
}
respMessage := c.server.HandleMessage(ctx, requestBytes)
respByte, err := json.Marshal(respMessage)
if err != nil {
return nil, fmt.Errorf("failed to marshal response message: %w", err)
}
var rpcResp JSONRPCResponse
err = json.Unmarshal(respByte, &rpcResp)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal response message: %w", err)
}
return &rpcResp, nil
}
func (c *InProcessTransport) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error {
notificationBytes, err := json.Marshal(notification)
if err != nil {
return fmt.Errorf("failed to marshal notification: %w", err)
}
notificationBytes = append(notificationBytes, '\n')
c.server.HandleMessage(ctx, notificationBytes)
return nil
}
func (c *InProcessTransport) SetNotificationHandler(handler func(notification mcp.JSONRPCNotification)) {
c.notifyMu.Lock()
defer c.notifyMu.Unlock()
c.onNotification = handler
}
func (c *InProcessTransport) Close() error {
if c.session != nil {
c.server.UnregisterSession(context.Background(), c.sessionID)
}
return nil
}
func (c *InProcessTransport) GetSessionId() string {
return ""
}
+74
View File
@@ -0,0 +1,74 @@
package transport
import (
"context"
"encoding/json"
"net/http"
"github.com/mark3labs/mcp-go/mcp"
)
// HTTPHeaderFunc is a function that extracts header entries from the given context
// and returns them as key-value pairs. This is typically used to add context values
// as HTTP headers in outgoing requests.
type HTTPHeaderFunc func(context.Context) map[string]string
// Interface for the transport layer.
type Interface interface {
// Start the connection. Start should only be called once.
Start(ctx context.Context) error
// SendRequest sends a json RPC request and returns the response synchronously.
SendRequest(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error)
// SendNotification sends a json RPC Notification to the server.
SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error
// SetNotificationHandler sets the handler for notifications.
// Any notification before the handler is set will be discarded.
SetNotificationHandler(handler func(notification mcp.JSONRPCNotification))
// Close the connection.
Close() error
// GetSessionId returns the session ID of the transport.
GetSessionId() string
}
// RequestHandler defines a function that handles incoming requests from the server.
type RequestHandler func(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error)
// BidirectionalInterface extends Interface to support incoming requests from the server.
// This is used for features like sampling where the server can send requests to the client.
type BidirectionalInterface interface {
Interface
// SetRequestHandler sets the handler for incoming requests from the server.
// The handler should process the request and return a response.
SetRequestHandler(handler RequestHandler)
}
// HTTPConnection is a Transport that runs over HTTP and supports
// protocol version headers.
type HTTPConnection interface {
Interface
SetProtocolVersion(version string)
}
type JSONRPCRequest struct {
JSONRPC string `json:"jsonrpc"`
ID mcp.RequestId `json:"id"`
Method string `json:"method"`
Params any `json:"params,omitempty"`
Header http.Header `json:"-"`
}
// JSONRPCResponse represents a JSON-RPC 2.0 response message.
// Use NewJSONRPCResultResponse to create a JSONRPCResponse with a result.
// Use NewJSONRPCErrorResponse to create a JSONRPCResponse with an error.
type JSONRPCResponse struct {
JSONRPC string `json:"jsonrpc"`
ID mcp.RequestId `json:"id"`
Result json.RawMessage `json:"result,omitempty"`
Error *mcp.JSONRPCErrorDetails `json:"error,omitempty"`
}
+712
View File
@@ -0,0 +1,712 @@
package transport
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"sync"
"time"
)
// ErrNoToken is returned when no token is available in the token store
var ErrNoToken = errors.New("no token available")
// OAuthConfig holds the OAuth configuration for the client
type OAuthConfig struct {
// ClientID is the OAuth client ID
ClientID string
// ClientSecret is the OAuth client secret (for confidential clients)
ClientSecret string
// RedirectURI is the redirect URI for the OAuth flow
RedirectURI string
// Scopes is the list of OAuth scopes to request
Scopes []string
// TokenStore is the storage for OAuth tokens
TokenStore TokenStore
// AuthServerMetadataURL is the URL to the OAuth server metadata
// If empty, the client will attempt to discover it from the base URL
AuthServerMetadataURL string
// PKCEEnabled enables PKCE for the OAuth flow (recommended for public clients)
PKCEEnabled bool
// HTTPClient is an optional HTTP client to use for requests.
// If nil, a default HTTP client with a 30 second timeout will be used.
HTTPClient *http.Client
}
// TokenStore is an interface for storing and retrieving OAuth tokens.
//
// Implementations must:
// - Honor context cancellation and deadlines, returning context.Canceled
// or context.DeadlineExceeded as appropriate
// - Return ErrNoToken (or a sentinel error that wraps it) when no token
// is available, rather than conflating this with other operational errors
// - Properly propagate all other errors (database failures, I/O errors, etc.)
// - Check ctx.Done() before performing operations and return ctx.Err() if cancelled
type TokenStore interface {
// GetToken returns the current token.
// Returns ErrNoToken if no token is available.
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
// Returns other errors for operational failures (I/O, database, etc.).
GetToken(ctx context.Context) (*Token, error)
// SaveToken saves a token.
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
// Returns other errors for operational failures (I/O, database, etc.).
SaveToken(ctx context.Context, token *Token) error
}
// Token represents an OAuth token
type Token struct {
// AccessToken is the OAuth access token
AccessToken string `json:"access_token"`
// TokenType is the type of token (usually "Bearer")
TokenType string `json:"token_type"`
// RefreshToken is the OAuth refresh token
RefreshToken string `json:"refresh_token,omitempty"`
// ExpiresIn is the number of seconds until the token expires
ExpiresIn int64 `json:"expires_in,omitempty"`
// Scope is the scope of the token
Scope string `json:"scope,omitempty"`
// ExpiresAt is the time when the token expires
ExpiresAt time.Time `json:"expires_at,omitempty"`
}
// IsExpired returns true if the token is expired
func (t *Token) IsExpired() bool {
if t.ExpiresAt.IsZero() {
return false
}
return time.Now().After(t.ExpiresAt)
}
// MemoryTokenStore is a simple in-memory token store
type MemoryTokenStore struct {
token *Token
mu sync.RWMutex
}
// NewMemoryTokenStore creates a new in-memory token store
func NewMemoryTokenStore() *MemoryTokenStore {
return &MemoryTokenStore{}
}
// GetToken returns the current token.
// Returns ErrNoToken if no token is available.
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
func (s *MemoryTokenStore) GetToken(ctx context.Context) (*Token, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
s.mu.RLock()
defer s.mu.RUnlock()
if s.token == nil {
return nil, ErrNoToken
}
return s.token, nil
}
// SaveToken saves a token.
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
func (s *MemoryTokenStore) SaveToken(ctx context.Context, token *Token) error {
if err := ctx.Err(); err != nil {
return err
}
s.mu.Lock()
defer s.mu.Unlock()
s.token = token
return nil
}
// AuthServerMetadata represents the OAuth 2.0 Authorization Server Metadata
type AuthServerMetadata struct {
Issuer string `json:"issuer"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
RegistrationEndpoint string `json:"registration_endpoint,omitempty"`
JwksURI string `json:"jwks_uri,omitempty"`
ScopesSupported []string `json:"scopes_supported,omitempty"`
ResponseTypesSupported []string `json:"response_types_supported"`
GrantTypesSupported []string `json:"grant_types_supported,omitempty"`
TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"`
}
// OAuthHandler handles OAuth authentication for HTTP requests
type OAuthHandler struct {
config OAuthConfig
httpClient *http.Client
serverMetadata *AuthServerMetadata
metadataFetchErr error
metadataOnce sync.Once
baseURL string
mu sync.RWMutex // Protects expectedState
expectedState string // Expected state value for CSRF protection
}
// NewOAuthHandler creates a new OAuth handler
func NewOAuthHandler(config OAuthConfig) *OAuthHandler {
if config.TokenStore == nil {
config.TokenStore = NewMemoryTokenStore()
}
if config.HTTPClient == nil {
config.HTTPClient = &http.Client{Timeout: 30 * time.Second}
}
return &OAuthHandler{
config: config,
httpClient: config.HTTPClient,
}
}
// GetAuthorizationHeader returns the Authorization header value for a request
func (h *OAuthHandler) GetAuthorizationHeader(ctx context.Context) (string, error) {
token, err := h.getValidToken(ctx)
if err != nil {
return "", err
}
// Some auth implementations are strict about token type
tokenType := token.TokenType
if tokenType == "bearer" {
tokenType = "Bearer"
}
return fmt.Sprintf("%s %s", tokenType, token.AccessToken), nil
}
// getValidToken returns a valid token, refreshing if necessary
func (h *OAuthHandler) getValidToken(ctx context.Context) (*Token, error) {
token, err := h.config.TokenStore.GetToken(ctx)
if err != nil && !errors.Is(err, ErrNoToken) {
return nil, err
}
if err == nil && !token.IsExpired() && token.AccessToken != "" {
return token, nil
}
// If we have a refresh token, try to use it
if err == nil && token.RefreshToken != "" {
newToken, err := h.refreshToken(ctx, token.RefreshToken)
if err == nil {
return newToken, nil
}
// If refresh fails, continue to authorization flow
}
// We need to get a new token through the authorization flow
return nil, ErrOAuthAuthorizationRequired
}
// refreshToken refreshes an OAuth token
func (h *OAuthHandler) refreshToken(ctx context.Context, refreshToken string) (*Token, error) {
metadata, err := h.getServerMetadata(ctx)
if err != nil {
return nil, fmt.Errorf("failed to get server metadata: %w", err)
}
data := url.Values{}
data.Set("grant_type", "refresh_token")
data.Set("refresh_token", refreshToken)
data.Set("client_id", h.config.ClientID)
if h.config.ClientSecret != "" {
data.Set("client_secret", h.config.ClientSecret)
}
req, err := http.NewRequestWithContext(
ctx,
http.MethodPost,
metadata.TokenEndpoint,
strings.NewReader(data.Encode()),
)
if err != nil {
return nil, fmt.Errorf("failed to create refresh token request: %w", err)
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
resp, err := h.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send refresh token request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, extractOAuthError(body, resp.StatusCode, "refresh token request failed")
}
var tokenResp Token
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
return nil, fmt.Errorf("failed to decode token response: %w", err)
}
// Set expiration time
if tokenResp.ExpiresIn > 0 {
tokenResp.ExpiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
}
// If no new refresh token is provided, keep the old one
if tokenResp.RefreshToken == "" {
tokenResp.RefreshToken = refreshToken
}
// Save the token
if err := h.config.TokenStore.SaveToken(ctx, &tokenResp); err != nil {
return nil, fmt.Errorf("failed to save token: %w", err)
}
return &tokenResp, nil
}
// RefreshToken is a public wrapper for refreshToken
func (h *OAuthHandler) RefreshToken(ctx context.Context, refreshToken string) (*Token, error) {
return h.refreshToken(ctx, refreshToken)
}
// GetClientID returns the client ID
func (h *OAuthHandler) GetClientID() string {
return h.config.ClientID
}
// extractOAuthError attempts to parse an OAuth error response from the response body
func extractOAuthError(body []byte, statusCode int, context string) error {
// Try to parse the error as an OAuth error response
var oauthErr OAuthError
if err := json.Unmarshal(body, &oauthErr); err == nil && oauthErr.ErrorCode != "" {
return fmt.Errorf("%s: %w", context, oauthErr)
}
// If not a valid OAuth error, return the raw response
return fmt.Errorf("%s with status %d: %s", context, statusCode, body)
}
// GetClientSecret returns the client secret
func (h *OAuthHandler) GetClientSecret() string {
return h.config.ClientSecret
}
// SetBaseURL sets the base URL for the API server
func (h *OAuthHandler) SetBaseURL(baseURL string) {
h.baseURL = baseURL
}
// GetExpectedState returns the expected state value (for testing purposes)
func (h *OAuthHandler) GetExpectedState() string {
h.mu.RLock()
defer h.mu.RUnlock()
return h.expectedState
}
// SetExpectedState sets the expected state value.
//
// This can be useful if you cannot maintain an OAuthHandler
// instance throughout the authentication flow; for example, if
// the initialization and callback steps are handled in different
// requests.
//
// In such cases, this should be called with the state value generated
// during the initial authentication request (e.g. by GenerateState)
// and included in the authorization URL.
func (h *OAuthHandler) SetExpectedState(expectedState string) {
h.mu.Lock()
defer h.mu.Unlock()
h.expectedState = expectedState
}
// OAuthError represents a standard OAuth 2.0 error response
type OAuthError struct {
ErrorCode string `json:"error"`
ErrorDescription string `json:"error_description,omitempty"`
ErrorURI string `json:"error_uri,omitempty"`
}
// Error implements the error interface
func (e OAuthError) Error() string {
if e.ErrorDescription != "" {
return fmt.Sprintf("OAuth error: %s - %s", e.ErrorCode, e.ErrorDescription)
}
return fmt.Sprintf("OAuth error: %s", e.ErrorCode)
}
// OAuthProtectedResource represents the response from /.well-known/oauth-protected-resource
type OAuthProtectedResource struct {
AuthorizationServers []string `json:"authorization_servers"`
Resource string `json:"resource"`
ResourceName string `json:"resource_name,omitempty"`
}
// getServerMetadata fetches the OAuth server metadata
func (h *OAuthHandler) getServerMetadata(ctx context.Context) (*AuthServerMetadata, error) {
h.metadataOnce.Do(func() {
// If AuthServerMetadataURL is explicitly provided, use it directly
if h.config.AuthServerMetadataURL != "" {
h.fetchMetadataFromURL(ctx, h.config.AuthServerMetadataURL)
return
}
// Try to discover the authorization server via OAuth Protected Resource
// as per RFC 9728 (https://datatracker.ietf.org/doc/html/rfc9728)
baseURL, err := h.extractBaseURL()
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to extract base URL: %w", err)
return
}
// Try to fetch the OAuth Protected Resource metadata
protectedResourceURL := baseURL + "/.well-known/oauth-protected-resource"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, protectedResourceURL, nil)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to create protected resource request: %w", err)
return
}
req.Header.Set("Accept", "application/json")
req.Header.Set("MCP-Protocol-Version", "2025-03-26")
resp, err := h.httpClient.Do(req)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to send protected resource request: %w", err)
return
}
defer resp.Body.Close()
// If we can't get the protected resource metadata, try OAuth Authorization Server discovery
if resp.StatusCode != http.StatusOK {
h.fetchMetadataFromURL(ctx, baseURL+"/.well-known/oauth-authorization-server")
if h.serverMetadata != nil {
return
}
// If that also fails, fall back to default endpoints
metadata, err := h.getDefaultEndpoints(baseURL)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
return
}
h.serverMetadata = metadata
return
}
// Parse the protected resource metadata
var protectedResource OAuthProtectedResource
if err := json.NewDecoder(resp.Body).Decode(&protectedResource); err != nil {
h.metadataFetchErr = fmt.Errorf("failed to decode protected resource response: %w", err)
return
}
// If no authorization servers are specified, fall back to default endpoints
if len(protectedResource.AuthorizationServers) == 0 {
metadata, err := h.getDefaultEndpoints(baseURL)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
return
}
h.serverMetadata = metadata
return
}
// Use the first authorization server
authServerURL := protectedResource.AuthorizationServers[0]
// Try OpenID Connect discovery first
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/openid-configuration")
if h.serverMetadata != nil {
return
}
// If OpenID Connect discovery fails, try OAuth Authorization Server Metadata
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/oauth-authorization-server")
if h.serverMetadata != nil {
return
}
// If both discovery methods fail, use default endpoints based on the authorization server URL
metadata, err := h.getDefaultEndpoints(authServerURL)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
return
}
h.serverMetadata = metadata
})
if h.metadataFetchErr != nil {
return nil, h.metadataFetchErr
}
return h.serverMetadata, nil
}
// fetchMetadataFromURL fetches and parses OAuth server metadata from a URL
func (h *OAuthHandler) fetchMetadataFromURL(ctx context.Context, metadataURL string) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to create metadata request: %w", err)
return
}
req.Header.Set("Accept", "application/json")
req.Header.Set("MCP-Protocol-Version", "2025-03-26")
resp, err := h.httpClient.Do(req)
if err != nil {
h.metadataFetchErr = fmt.Errorf("failed to send metadata request: %w", err)
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
// If metadata discovery fails, don't set any metadata
return
}
var metadata AuthServerMetadata
if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
h.metadataFetchErr = fmt.Errorf("failed to decode metadata response: %w", err)
return
}
h.serverMetadata = &metadata
}
// extractBaseURL extracts the base URL from the first request
func (h *OAuthHandler) extractBaseURL() (string, error) {
// If we have a base URL from a previous request, use it
if h.baseURL != "" {
return h.baseURL, nil
}
// Otherwise, we need to infer it from the redirect URI
if h.config.RedirectURI == "" {
return "", fmt.Errorf("no base URL available and no redirect URI provided")
}
// Parse the redirect URI to extract the authority
parsedURL, err := url.Parse(h.config.RedirectURI)
if err != nil {
return "", fmt.Errorf("failed to parse redirect URI: %w", err)
}
// Use the scheme and host from the redirect URI
baseURL := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host)
return baseURL, nil
}
// GetServerMetadata is a public wrapper for getServerMetadata
func (h *OAuthHandler) GetServerMetadata(ctx context.Context) (*AuthServerMetadata, error) {
return h.getServerMetadata(ctx)
}
// getDefaultEndpoints returns default OAuth endpoints based on the base URL
func (h *OAuthHandler) getDefaultEndpoints(baseURL string) (*AuthServerMetadata, error) {
// Parse the base URL to extract the authority
parsedURL, err := url.Parse(baseURL)
if err != nil {
return nil, fmt.Errorf("failed to parse base URL: %w", err)
}
// Discard any path component to get the authorization base URL
parsedURL.Path = ""
authBaseURL := parsedURL.String()
// Validate that the URL has a scheme and host
if parsedURL.Scheme == "" || parsedURL.Host == "" {
return nil, fmt.Errorf("invalid base URL: missing scheme or host in %q", baseURL)
}
return &AuthServerMetadata{
Issuer: authBaseURL,
AuthorizationEndpoint: authBaseURL + "/authorize",
TokenEndpoint: authBaseURL + "/token",
RegistrationEndpoint: authBaseURL + "/register",
}, nil
}
// RegisterClient performs dynamic client registration
func (h *OAuthHandler) RegisterClient(ctx context.Context, clientName string) error {
metadata, err := h.getServerMetadata(ctx)
if err != nil {
return fmt.Errorf("failed to get server metadata: %w", err)
}
if metadata.RegistrationEndpoint == "" {
return errors.New("server does not support dynamic client registration")
}
// Prepare registration request
regRequest := map[string]any{
"client_name": clientName,
"redirect_uris": []string{h.config.RedirectURI},
"token_endpoint_auth_method": "none", // For public clients
"grant_types": []string{"authorization_code", "refresh_token"},
"response_types": []string{"code"},
"scope": strings.Join(h.config.Scopes, " "),
}
// Add client_secret if this is a confidential client
if h.config.ClientSecret != "" {
regRequest["token_endpoint_auth_method"] = "client_secret_basic"
}
reqBody, err := json.Marshal(regRequest)
if err != nil {
return fmt.Errorf("failed to marshal registration request: %w", err)
}
req, err := http.NewRequestWithContext(
ctx,
http.MethodPost,
metadata.RegistrationEndpoint,
bytes.NewReader(reqBody),
)
if err != nil {
return fmt.Errorf("failed to create registration request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
resp, err := h.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send registration request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return extractOAuthError(body, resp.StatusCode, "registration request failed")
}
var regResponse struct {
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret,omitempty"`
}
if err := json.NewDecoder(resp.Body).Decode(&regResponse); err != nil {
return fmt.Errorf("failed to decode registration response: %w", err)
}
// Update the client configuration
h.config.ClientID = regResponse.ClientID
if regResponse.ClientSecret != "" {
h.config.ClientSecret = regResponse.ClientSecret
}
return nil
}
// ErrInvalidState is returned when the state parameter doesn't match the expected value
var ErrInvalidState = errors.New("invalid state parameter, possible CSRF attack")
// ProcessAuthorizationResponse processes the authorization response and exchanges the code for a token
func (h *OAuthHandler) ProcessAuthorizationResponse(ctx context.Context, code, state, codeVerifier string) error {
// Validate the state parameter to prevent CSRF attacks
h.mu.Lock()
expectedState := h.expectedState
if expectedState == "" {
h.mu.Unlock()
return errors.New("no expected state found, authorization flow may not have been initiated properly")
}
if state != expectedState {
h.mu.Unlock()
return ErrInvalidState
}
// Clear the expected state after validation
h.expectedState = ""
h.mu.Unlock()
metadata, err := h.getServerMetadata(ctx)
if err != nil {
return fmt.Errorf("failed to get server metadata: %w", err)
}
data := url.Values{}
data.Set("grant_type", "authorization_code")
data.Set("code", code)
data.Set("client_id", h.config.ClientID)
data.Set("redirect_uri", h.config.RedirectURI)
if h.config.ClientSecret != "" {
data.Set("client_secret", h.config.ClientSecret)
}
if h.config.PKCEEnabled && codeVerifier != "" {
data.Set("code_verifier", codeVerifier)
}
req, err := http.NewRequestWithContext(
ctx,
http.MethodPost,
metadata.TokenEndpoint,
strings.NewReader(data.Encode()),
)
if err != nil {
return fmt.Errorf("failed to create token request: %w", err)
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
resp, err := h.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send token request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return extractOAuthError(body, resp.StatusCode, "token request failed")
}
var tokenResp Token
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
return fmt.Errorf("failed to decode token response: %w", err)
}
// Set expiration time
if tokenResp.ExpiresIn > 0 {
tokenResp.ExpiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
}
// Save the token
if err := h.config.TokenStore.SaveToken(ctx, &tokenResp); err != nil {
return fmt.Errorf("failed to save token: %w", err)
}
return nil
}
// GetAuthorizationURL returns the URL for the authorization endpoint
func (h *OAuthHandler) GetAuthorizationURL(ctx context.Context, state, codeChallenge string) (string, error) {
metadata, err := h.getServerMetadata(ctx)
if err != nil {
return "", fmt.Errorf("failed to get server metadata: %w", err)
}
// Store the state for later validation
h.SetExpectedState(state)
params := url.Values{}
params.Set("response_type", "code")
params.Set("client_id", h.config.ClientID)
params.Set("redirect_uri", h.config.RedirectURI)
params.Set("state", state)
if len(h.config.Scopes) > 0 {
params.Set("scope", strings.Join(h.config.Scopes, " "))
}
if h.config.PKCEEnabled && codeChallenge != "" {
params.Set("code_challenge", codeChallenge)
params.Set("code_challenge_method", "S256")
}
return metadata.AuthorizationEndpoint + "?" + params.Encode(), nil
}
Loaded 100 of 190 files, more files were not shown because too many files have changed in this diff. Show more