Skip to content

Commit

Permalink
Use an existing logger for sampling:
Browse files Browse the repository at this point in the history
This library allows for sampling log messages
that match some certain criteria. I have not implemented
this in Smee, but this could allow for sampling each
unique mac address in the url path. This would be as
opposed to having all traffic be sampled together.
This needs tested to see if it's worth it. If it's
not there are other ways to do the sampling in the same
way without needing to import a library.

Signed-off-by: Jacob Weinstock <[email protected]>
  • Loading branch information
jacobweinstock committed Nov 18, 2024
1 parent b1a06b1 commit 753dc2d
Show file tree
Hide file tree
Showing 5 changed files with 107 additions and 59 deletions.
15 changes: 15 additions & 0 deletions cmd/smee/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"errors"
"flag"
"fmt"
"log/slog"
"net"
"net/netip"
"net/url"
Expand All @@ -19,6 +20,8 @@ import (
"github.com/go-logr/zapr"
"github.com/insomniacslk/dhcp/dhcpv4"
"github.com/insomniacslk/dhcp/dhcpv4/server4"
slogmulti "github.com/samber/slog-multi"
slogsampling "github.com/samber/slog-sampling"
"github.com/tinkerbell/ipxedust"
"github.com/tinkerbell/ipxedust/ihttp"
"github.com/tinkerbell/smee/internal/dhcp/handler"
Expand Down Expand Up @@ -255,6 +258,17 @@ func main() {
if err != nil {
panic(fmt.Errorf("failed to create backend: %w", err))
}
// Will print 10% of entries.
option := slogsampling.UniformSamplingOption{
// The sample rate for sampling traces in the range [0.0, 1.0].
Rate: 0.002,
}

logger := slog.New(
slogmulti.
Pipe(option.NewMiddleware()).
Handler(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{AddSource: true})),
)
ih := iso.Handler{
Logger: log,
Backend: br,
Expand All @@ -269,6 +283,7 @@ func main() {
}
return cfg.iso.magicString
}(),
SampleLogger: logger,
}
isoHandler, err := ih.Reverse()
if err != nil {
Expand Down
6 changes: 6 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ require (
github.com/packethost/xff v0.0.0-20190305172552-d3e9190c41b3
github.com/peterbourgon/ff/v3 v3.4.0
github.com/prometheus/client_golang v1.20.5
github.com/samber/slog-multi v1.2.4
github.com/samber/slog-sampling v1.5.2
github.com/tinkerbell/ipxedust v0.0.0-20241108174245-aa0c0298057d
github.com/tinkerbell/tink v0.12.1
github.com/vishvananda/netlink v1.3.0
Expand All @@ -37,8 +39,10 @@ require (
require (
dario.cat/mergo v1.0.1 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/bluele/gcache v0.0.2 // indirect
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cornelk/hashmap v1.0.8 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/emicklei/go-restful/v3 v3.12.1 // indirect
github.com/evanphx/json-patch/v5 v5.9.0 // indirect
Expand Down Expand Up @@ -80,6 +84,8 @@ require (
github.com/prometheus/common v0.55.0 // indirect
github.com/prometheus/procfs v0.15.1 // indirect
github.com/rs/zerolog v1.33.0 // indirect
github.com/samber/lo v1.47.0 // indirect
github.com/samber/slog-common v0.17.1 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/u-root/uio v0.0.0-20230305220412-3e8cd9d6bf63 // indirect
github.com/vishvananda/netns v0.0.4 // indirect
Expand Down
12 changes: 12 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,17 @@ dario.cat/mergo v1.0.1 h1:Ra4+bf83h2ztPIQYNP99R6m+Y7KfnARDfID+a+vLl4s=
dario.cat/mergo v1.0.1/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk=
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
github.com/bluele/gcache v0.0.2 h1:WcbfdXICg7G/DGBh1PFfcirkWOQV+v077yF1pSy3DGw=
github.com/bluele/gcache v0.0.2/go.mod h1:m15KV+ECjptwSPxKhOhQoAFQVtUFjTVkc3H8o0t/fp0=
github.com/ccoveille/go-safecast v1.1.0 h1:iHKNWaZm+OznO7Eh6EljXPjGfGQsSfa6/sxPlIEKO+g=
github.com/ccoveille/go-safecast v1.1.0/go.mod h1:QqwNjxQ7DAqY0C721OIO9InMk9zCwcsO7tnRuHytad8=
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/cornelk/hashmap v1.0.8 h1:nv0AWgw02n+iDcawr5It4CjQIAcdMMKRrs10HOJYlrc=
github.com/cornelk/hashmap v1.0.8/go.mod h1:RfZb7JO3RviW/rT6emczVuC/oxpdz4UsSB2LJSclR1k=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
Expand Down Expand Up @@ -150,6 +154,14 @@ github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWN
github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
github.com/samber/slog-common v0.17.1 h1:jTqqLBgoJshpoxlPSGiypyOanjH6tY+i9bwyYmIbjhI=
github.com/samber/slog-common v0.17.1/go.mod h1:mZSJhinB4aqHziR0SKPqpVZjJ0JO35JfH+dDIWqaCBk=
github.com/samber/slog-multi v1.2.4 h1:k9x3JAWKJFPKffx+oXZ8TasaNuorIW4tG+TXxkt6Ry4=
github.com/samber/slog-multi v1.2.4/go.mod h1:ACuZ5B6heK57TfMVkVknN2UZHoFfjCwRxR0Q2OXKHlo=
github.com/samber/slog-sampling v1.5.2 h1:HaQmRGmLkVsXlVHEgIzgvxc6JYkkppL/7GN1B1g5LQM=
github.com/samber/slog-sampling v1.5.2/go.mod h1:n2PVbLAFRx8Im0KFt9srRrO/8wi9FAiTNBV7FmhbAHY=
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
Expand Down
3 changes: 3 additions & 0 deletions internal/ipxe/http/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ func (h *loggingMiddleware) ServeHTTP(w http.ResponseWriter, req *http.Request)
res := &responseWriter{ResponseWriter: w}
h.handler.ServeHTTP(res, req) // process the request

// The "X-Global-Logging" header allows all registered HTTP handlers to disable this global logging
// by setting the header to any non empty string. This is useful for handlers that handle partial content of
// larger file. The ISO handler, for example.
r := res.Header().Get("X-Global-Logging")

if log && r == "" {
Expand Down
130 changes: 71 additions & 59 deletions internal/iso/iso.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@ import (
"errors"
"fmt"
"io"
"log/slog"
"net"
"net/http"
"net/http/httputil"
"net/url"
"path"
"path/filepath"
"strings"
"sync/atomic"

apierrors "k8s.io/apimachinery/pkg/api/errors"

Expand All @@ -21,7 +23,7 @@ import (
)

const (
defaultConsoles = "console=ttyS1 console=ttyS1 console=ttyS0 console=ttyAMA0 console=ttyS1 console=tty0"
defaultConsoles = "console=ttyS0 console=ttyAMA0 console=tty1 console=tty0, console=ttyS1,115200"
)

type Handler struct {
Expand All @@ -41,12 +43,17 @@ type Handler struct {
// in the source iso before patching. The field can be set
// during build time by setting this field.
// Ref: https://github.com/tinkerbell/hook/blob/main/linuxkit-templates/hook.template.yaml
MagicString string
MagicString string
SampleLogger *slog.Logger
}

var total atomic.Int64

func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {
log := h.Logger.WithValues("method", req.Method, "url", req.URL.Path, "xff", req.Header.Get("X-Forwarded-For"), "remoteAddr", req.RemoteAddr)
log.V(1).Info("starting the patching function")
log := h.Logger.WithValues("method", req.Method, "urlPath", req.URL.Path, "xff", req.Header.Get("X-Forwarded-For"), "remoteAddr", req.RemoteAddr, "total", total.Load())
splog := h.SampleLogger.With("method", req.Method, "urlPath", req.URL.Path, "xff", req.Header.Get("X-Forwarded-For"), "remoteAddr", req.RemoteAddr, "total", total.Load())
total.Add(1)
log.V(1).Info("starting the ISO patching HTTP handler")
if req.Method != http.MethodHead && req.Method != http.MethodGet {
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusNotImplemented, http.StatusText(http.StatusNotImplemented)),
Expand All @@ -57,7 +64,7 @@ func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {
}

if filepath.Ext(req.URL.Path) != ".iso" {
log.Info("Extension not supported, only supported type is '.iso'", "path", req.URL.Path)
log.Info("extension not supported, only supported extension is '.iso'")
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusNotFound, http.StatusText(http.StatusNotFound)),
StatusCode: http.StatusNotFound,
Expand All @@ -73,9 +80,9 @@ func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {
// We serve the iso only if this validation passes.
ha, err := getMAC(req.URL.Path)
if err != nil {
log.Info("unable to get the mac address", "error", err)
log.Info("unable to parse mac address in the URL path", "error", err)
return &http.Response{
Status: "400 BAD REQUEST",
Status: fmt.Sprintf("%d %s", http.StatusBadRequest, http.StatusText(http.StatusBadRequest)),
StatusCode: http.StatusBadRequest,
Body: http.NoBody,
Request: req,
Expand All @@ -84,7 +91,7 @@ func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {

f, err := getFacility(req.Context(), ha, h.Backend)
if err != nil {
log.Info("unable to get facility", "mac", ha, "error", err)
log.Error(err, "unable to get the hardware object", "mac", ha)
if apierrors.IsNotFound(err) {
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusNotFound, http.StatusText(http.StatusNotFound)),
Expand All @@ -111,7 +118,7 @@ func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {
// For our use case the default transport will suffice.
resp, err := http.DefaultTransport.RoundTrip(req)
if err != nil {
log.Info("issue with getting the source ISO", "sourceIso", h.SourceISO, "error", err)
log.Info("issue getting the source ISO", "sourceIso", h.SourceISO, "error", err)
return nil, err
}
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
Expand All @@ -123,63 +130,68 @@ func (h *Handler) RoundTrip(req *http.Request) (*http.Response, error) {
resp.Header.Set("X-Global-Logging", "false")

if req.Method == http.MethodHead {
// Fuse client typically make a HEAD request before they start requesting content.
log.V(1).Info("HTTP HEAD request received, patching only occurs on 206 requests")
// Fuse clients typically make a HEAD request before they start requesting content.
log.Info("HTTP HEAD method received")
return resp, nil
}

// roundtripper should only return error when no response from the server
// for any other case just log the error and return 404 response
if resp.StatusCode == http.StatusPartialContent {
b, err := io.ReadAll(resp.Body)
if err != nil {
log.Error(err, "reading response bytes")
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError)),
StatusCode: http.StatusInternalServerError,
Body: http.NoBody,
Request: req,
}, nil
}
if err := resp.Body.Close(); err != nil {
log.Error(err, "closing response body")
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError)),
StatusCode: http.StatusInternalServerError,
Body: http.NoBody,
Request: req,
}, nil
}
if resp.StatusCode == http.StatusOK {
log.Info("HTTP GET method received with a 200 status code")
return resp, nil
}

// The hardware object doesn't contain a specific field for consoles
// right now facility is used instead.
var consoles string
switch {
case f != "" && strings.Contains(f, "console="):
consoles = fmt.Sprintf("facility=%s", f)
case f != "":
consoles = fmt.Sprintf("facility=%s %s", f, defaultConsoles)
default:
consoles = defaultConsoles
}
magicStringPadding := bytes.Repeat([]byte{' '}, len(h.MagicString))

// TODO: revisit later to handle the magic string potentially being spread across two chunks.
// In current implementation we will never patch the above case. Add logic to patch the case of
// magic string spread across multiple response bodies in the future.
i := bytes.Index(b, []byte(h.MagicString))
if i != -1 {
log.Info("magic string found, patching the ISO")
dup := make([]byte, len(b))
copy(dup, b)
copy(dup[i:], magicStringPadding)
copy(dup[i:], []byte(h.constructPatch(consoles, ha.String())))
b = dup
}
splog.Info("HTTP GET method received with a 206 status code")

resp.Body = io.NopCloser(bytes.NewReader(b))
// this roundtripper func should only return error when there is no response from the server.
// for any other case we log the error and return a 500 response
b, err := io.ReadAll(resp.Body)
if err != nil {
log.Error(err, "issue reading response bytes")
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError)),
StatusCode: http.StatusInternalServerError,
Body: http.NoBody,
Request: req,
}, nil
}
if err := resp.Body.Close(); err != nil {
log.Error(err, "issue closing response body")
return &http.Response{
Status: fmt.Sprintf("%d %s", http.StatusInternalServerError, http.StatusText(http.StatusInternalServerError)),
StatusCode: http.StatusInternalServerError,
Body: http.NoBody,
Request: req,
}, nil
}

// The hardware object doesn't contain a dedicated field for consoles right now and
// historically the facility is used as a way to define consoles on a per Hardware basis.
var consoles string
switch {
case f != "" && strings.Contains(f, "console="):
consoles = fmt.Sprintf("facility=%s", f)
case f != "":
consoles = fmt.Sprintf("facility=%s %s", f, defaultConsoles)
default:
consoles = defaultConsoles
}
magicStringPadding := bytes.Repeat([]byte{' '}, len(h.MagicString))

// TODO: revisit later to handle the magic string potentially being spread across two chunks.
// In current implementation we will never patch the above case. Add logic to patch the case of
// magic string spread across multiple response bodies in the future.
i := bytes.Index(b, []byte(h.MagicString))
if i != -1 {
log.Info("magic string found, patching the content", "contentRange", resp.Header.Get("Content-Range"))
dup := make([]byte, len(b))
copy(dup, b)
copy(dup[i:], magicStringPadding)
copy(dup[i:], []byte(h.constructPatch(consoles, ha.String())))
b = dup
}

resp.Body = io.NopCloser(bytes.NewReader(b))

log.V(1).Info("roundtrip complete")
return resp, nil
}
Expand Down

0 comments on commit 753dc2d

Please sign in to comment.