1
4
mirror of https://github.com/yunginnanet/HellPot synced 2024-06-20 12:58:02 +00:00

Compare commits

...

58 Commits

Author SHA1 Message Date
478ebe3162
Merge 7e100317ab5c96eb8baa6e6cbb45911f235ba312 into 8146ee46edc19665cd9ae0369a8f106c65a6596f 2023-11-16 07:39:53 +00:00
7e100317ab
Chore[CD]: Deprecate unused workflow 2023-11-15 23:39:45 -08:00
c717e6ec5c
Merge branch 'main' into development 2023-11-15 23:38:03 -08:00
8146ee46ed
Update README.md, add syntax highlighting 2023-11-15 23:37:40 -08:00
dependabot[bot]
dadb8dd8ee Bump git.tcp.direct/kayos/common from 0.9.4 to 0.9.6
Bumps git.tcp.direct/kayos/common from 0.9.4 to 0.9.6.

---
updated-dependencies:
- dependency-name: git.tcp.direct/kayos/common
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-11-15 23:35:16 -08:00
f80435d2e4
Merge branch 'main' into development 2023-11-15 23:34:54 -08:00
955e49e172
Fix[CI]: fix workflow branch name
(cherry picked from commit c58a3465aef6a56adfa23b76b4313f26997be4e3)
Signed-off-by: kayos@tcp.direct <kayos@tcp.direct>
2023-11-15 23:33:18 -08:00
c58a3465ae
Fix[CI]: fix workflow branch name 2023-11-15 23:32:08 -08:00
dependabot[bot]
6f732e71a4 Bump github.com/valyala/fasthttp from 1.50.0 to 1.51.0
Bumps [github.com/valyala/fasthttp](https://github.com/valyala/fasthttp) from 1.50.0 to 1.51.0.
- [Release notes](https://github.com/valyala/fasthttp/releases)
- [Commits](https://github.com/valyala/fasthttp/compare/v1.50.0...v1.51.0)

---
updated-dependencies:
- dependency-name: github.com/valyala/fasthttp
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-11-14 14:42:52 -08:00
add6bd55d7
Fix[CI]: gosec nosec (again) 2023-10-20 01:17:21 -07:00
3c39720646
Chore: tidy up 2023-10-20 01:04:55 -07:00
3f79f26809
Merge branch 'main' into development 2023-10-19 23:18:27 -07:00
7633aa3e79
Fix: remediate unit test race condition 2023-10-19 04:14:53 -07:00
7c5cc69038
Chore: tidy up 2023-10-19 04:09:17 -07:00
e3e05258e7
Fix[CI](SAST): gosec nosec 2023-10-19 04:07:52 -07:00
d88ed900f8
Feat[speedo]: flesh out speedo and add TCP transfer test 2023-10-19 03:52:08 -07:00
ec44773c45
Chore: deps 2023-08-11 23:07:38 -07:00
dependabot[bot]
37ce853d09
Bump git.tcp.direct/kayos/common from 0.8.5 to 0.8.6
Bumps git.tcp.direct/kayos/common from 0.8.5 to 0.8.6.

---
updated-dependencies:
- dependency-name: git.tcp.direct/kayos/common
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-08-11 23:07:38 -07:00
dependabot[bot]
85ca012599
Bump git.tcp.direct/kayos/common from 0.8.4 to 0.8.5
Bumps git.tcp.direct/kayos/common from 0.8.4 to 0.8.5.

---
updated-dependencies:
- dependency-name: git.tcp.direct/kayos/common
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-08-11 23:07:38 -07:00
dependabot[bot]
da957a4a5d
Bump github.com/fasthttp/router from 1.4.18 to 1.4.19
Bumps [github.com/fasthttp/router](https://github.com/fasthttp/router) from 1.4.18 to 1.4.19.
- [Release notes](https://github.com/fasthttp/router/releases)
- [Commits](https://github.com/fasthttp/router/compare/v1.4.18...v1.4.19)

---
updated-dependencies:
- dependency-name: github.com/fasthttp/router
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-08-11 23:07:38 -07:00
dependabot[bot]
42eea5f9dd
Bump git.tcp.direct/kayos/common from 0.8.3 to 0.8.4
Bumps git.tcp.direct/kayos/common from 0.8.3 to 0.8.4.

---
updated-dependencies:
- dependency-name: git.tcp.direct/kayos/common
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
0005356845
Bump golang.org/x/term from 0.7.0 to 0.8.0
Bumps [golang.org/x/term](https://github.com/golang/term) from 0.7.0 to 0.8.0.
- [Commits](https://github.com/golang/term/compare/v0.7.0...v0.8.0)

---
updated-dependencies:
- dependency-name: golang.org/x/term
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
3cd2b7a862
Bump git.tcp.direct/kayos/common from 0.8.2 to 0.8.3 (#87)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
698b4bdda9
Bump github.com/valyala/fasthttp from 1.46.0 to 1.47.0 (#85)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: kayos <kayos@tcp.direct>
2023-08-11 23:07:37 -07:00
dependabot[bot]
57526ab3a1
Bump wangyoucao577/go-release-action from 1.37 to 1.38 (#86)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
bab07d7f92
Bump github.com/valyala/fasthttp from 1.45.0 to 1.46.0 (#84)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
2333643040
Bump github.com/rs/zerolog from 1.29.0 to 1.29.1 (#83)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:37 -07:00
dependabot[bot]
3776301ed0
Bump golang.org/x/term from 0.6.0 to 0.7.0 (#81)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
dependabot[bot]
9ca252dad4
Bump peter-evans/create-or-update-comment from 2 to 3 (#82)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
dependabot[bot]
a635c6d252
Bump actions/checkout from 2 to 3 (#79)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
dependabot[bot]
bca0f4491a
Bump wangyoucao577/go-release-action from 1.35 to 1.37 (#80)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
dependabot[bot]
ba8ee74863
Bump actions/setup-go from 2 to 4 (#78)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
2bfe90cadc
Update dependabot.yml 2023-08-11 23:07:36 -07:00
dependabot[bot]
03113b443a
Bump git.tcp.direct/kayos/common from 0.8.1 to 0.8.2 (#77)
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-08-11 23:07:36 -07:00
4126e2c5ee
Merge branch 'master' into development 2023-03-23 01:26:01 -07:00
34323c67f6
Fix: HTTP server failing to listen should be fatal 2023-03-23 01:22:08 -07:00
6d43e2e6d2
CI: Add race detector 2023-03-05 03:56:37 -08:00
c61f6b4a9c
Feat: Fix and finish bandwidth limiting io.Writer 2023-03-05 03:56:26 -08:00
b8b4b56cba
Fix: Re-add speedometer for bandwidth metering feature 2023-03-05 01:13:07 -08:00
e99b81b34c
Merge branch 'master' into development 2023-03-05 01:07:18 -08:00
bcd514c95f
CD: Fix last commit 2023-02-02 21:22:51 -08:00
85897a2e2c
CD: Enable automatic builds and releases 2023-02-02 21:18:47 -08:00
098e21e803
Merge branch 'master' into development 2023-02-02 20:41:25 -08:00
141669f41a
Merge branch 'master' into development 2023-02-02 20:40:35 -08:00
6d00df0792
Merge branch 'master' into development 2023-01-03 09:50:49 -08:00
ecfcc06823
Lint/refactor 2022-12-31 17:11:38 -08:00
2210243788
Merge branch 'master' into development 2022-12-17 20:35:29 -08:00
7c0edb9708
CI: Update 2022-09-11 04:36:26 -07:00
0461e01a60
Fix: log directory 2022-09-11 04:30:27 -07:00
20e54b74ff
Merge branch 'master' into development 2022-09-11 04:27:21 -07:00
192e18cd54
Fix: Dev versioning 2022-09-11 04:25:11 -07:00
c86b09018e
Selective merge from development branch
Signed-off-by: kayos@tcp.direct <kayos@tcp.direct>
2022-09-11 04:07:33 -07:00
b2db55a9d5
Update development versioning mechanism 2022-09-11 04:06:00 -07:00
9094317099
Lint: dead code removal 2022-09-11 03:49:41 -07:00
84b44f5db0
Adjust robots.txt generation 2022-09-11 03:49:40 -07:00
64b0373e52
Minor adjustments 2022-09-11 03:49:37 -07:00
dca759ac37
Aesthetic: adjust banner version handling 2022-09-11 03:47:01 -07:00
01e4853475
Tidy up: stop using globals 2022-09-11 03:46:55 -07:00
8 changed files with 654 additions and 20 deletions

@ -2,9 +2,9 @@ name: Vibe Check
on:
push:
branches: [ master, development ]
branches: [ main, development ]
pull_request:
branches: [ master ]
branches: [ main ]
jobs:
build:
@ -14,7 +14,7 @@ jobs:
- name: Set up Go
uses: actions/setup-go@v4
with:
go-version: 1.19
go-version: 1.21
- name: go vet -v ./...
run: go vet -v ./...
- name: gosec ./...
@ -22,7 +22,7 @@ jobs:
export PATH=$PATH:$(go env GOPATH)/bin
go install github.com/securego/gosec/v2/cmd/gosec@latest
gosec ./...
- name: go test -v ./...
run: go test -v ./...
- name: go test -race -v ./...
run: go test -race -v ./...
- name: go build -v ./...
run: go build -v ./...

@ -111,7 +111,7 @@ In the event of a missing configuration file, HellPot will attempt to place it's
## Example Web Server Config (nginx)
```
```nginx
location '/robots.txt' {
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
@ -133,7 +133,7 @@ All nonexisting URLs are being reverse proxied to a HellPot instance on localhos
* Requests on nonexisting URLs cause a HTTP Error 404, which content is served by HellPot
* URLs under the "/.well-known/" suffix are excluded.
```
```apache
<VirtualHost yourserver>
ErrorDocument 400 "/content/400"
ErrorDocument 403 "/content/403"

@ -50,7 +50,7 @@ func main() {
signal.Notify(stopChan, syscall.SIGINT, syscall.SIGTERM)
go func() {
log.Error().Err(http.Serve()).Msg("HTTP error")
log.Fatal().Err(http.Serve()).Msg("HTTP error")
}()
<-stopChan // wait for SIGINT

4
go.mod

@ -3,12 +3,12 @@ module github.com/yunginnanet/HellPot
go 1.19
require (
git.tcp.direct/kayos/common v0.9.4
git.tcp.direct/kayos/common v0.9.6
github.com/fasthttp/router v1.4.21
github.com/rs/zerolog v1.31.0
github.com/spf13/afero v1.10.0
github.com/spf13/viper v1.17.0
github.com/valyala/fasthttp v1.50.0
github.com/valyala/fasthttp v1.51.0
golang.org/x/term v0.14.0
)

8
go.sum

@ -36,8 +36,8 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX
cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0=
cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo=
dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
git.tcp.direct/kayos/common v0.9.4 h1:UqenO7AWOeDuItziNWQNyoBbWRDPO2GofNoYcTF60lc=
git.tcp.direct/kayos/common v0.9.4/go.mod h1:tTqUGj50mpwoQD0Zsdsv6cpDzN9VfjnQMgpDC8aRfms=
git.tcp.direct/kayos/common v0.9.6 h1:EITtktxZF/zkzqAhZZxvm6cZpFYoZ0P/gLB9RPatKUY=
git.tcp.direct/kayos/common v0.9.6/go.mod h1:8y9b+PN1+ZVaQ/VugD9dkKe+uqhE8jH7a64RyF7h2rM=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs=
@ -193,8 +193,8 @@ github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasthttp v1.50.0 h1:H7fweIlBm0rXLs2q0XbalvJ6r0CUPFWK3/bB4N13e9M=
github.com/valyala/fasthttp v1.50.0/go.mod h1:k2zXd82h/7UZc3VOdJ2WaUqt1uZ/XpXAfE9i+HBC3lA=
github.com/valyala/fasthttp v1.51.0 h1:8b30A5JlZ6C7AS81RsWjYMQmrZG6feChmgAolCl1SqA=
github.com/valyala/fasthttp v1.51.0/go.mod h1:oI2XroL+lI7vdXyYoQk03bXBThfFl2cVdIA3Xl7cH8g=
github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

@ -18,9 +18,7 @@ var (
logger zerolog.Logger
)
// StartLogger instantiates an instance of our zerolog loggger so we can hook it in our main package.
// While this does return a logger, it should not be used for additional retrievals of the logger. Use GetLogger()
func StartLogger(pretty bool, targets ...io.Writer) zerolog.Logger {
func prepLogDir() {
logDir = snek.GetString("logger.directory")
if !strings.HasSuffix(logDir, "/") {
logDir += "/"
@ -29,7 +27,11 @@ func StartLogger(pretty bool, targets ...io.Writer) zerolog.Logger {
println("cannot create log directory: " + logDir + "(" + err.Error() + ")")
os.Exit(1)
}
}
// StartLogger instantiates an instance of our zerolog loggger so we can hook it in our main package.
// While this does return a logger, it should not be used for additional retrievals of the logger. Use GetLogger().
func StartLogger(pretty bool, targets ...io.Writer) zerolog.Logger {
logFileName := "HellPot"
if snek.GetBool("logger.use_date_filename") {
@ -44,9 +46,11 @@ func StartLogger(pretty bool, targets ...io.Writer) zerolog.Logger {
case len(targets) > 0:
logFile = io.MultiWriter(targets...)
default:
prepLogDir()
CurrentLogFile = path.Join(logDir, logFileName+".log")
/* #nosec */
if logFile, err = os.OpenFile(CurrentLogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o666); err != nil {
//nolint:lll
logFile, err = os.OpenFile(CurrentLogFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o666) // #nosec G304 G302 -- we are not using user input to create the file
if err != nil {
println("cannot create log file: " + err.Error())
os.Exit(1)
}
@ -62,7 +66,7 @@ func StartLogger(pretty bool, targets ...io.Writer) zerolog.Logger {
return logger
}
// GetLogger retrieves our global logger object
// GetLogger retrieves our global logger object.
func GetLogger() *zerolog.Logger {
// future logic here
return &logger

@ -0,0 +1,237 @@
package util
import (
"errors"
"fmt"
"io"
"sync"
"sync/atomic"
"time"
)
var ErrLimitReached = errors.New("limit reached")
// Speedometer is an io.Writer wrapper that will limit the rate at which data is written to the underlying target.
//
// It is safe for concurrent use, but writers will block when slowed down.
//
// Optionally, it can be given;
//
// - a capacity, which will cause it to return an error if the capacity is exceeded.
//
// - a speed limit, causing slow downs of data written to the underlying writer if the speed limit is exceeded.
type Speedometer struct {
ceiling int64
speedLimit *SpeedLimit
internal atomics
w io.Writer
}
type atomics struct {
count *atomic.Int64
closed *atomic.Bool
start *sync.Once
stop *sync.Once
birth *atomic.Pointer[time.Time]
duration *atomic.Pointer[time.Duration]
slow *atomic.Bool
}
func newAtomics() atomics {
manhattan := atomics{
count: new(atomic.Int64),
closed: new(atomic.Bool),
start: new(sync.Once),
stop: new(sync.Once),
birth: new(atomic.Pointer[time.Time]),
duration: new(atomic.Pointer[time.Duration]),
slow: new(atomic.Bool),
}
manhattan.birth.Store(&time.Time{})
manhattan.closed.Store(false)
manhattan.count.Store(0)
return manhattan
}
// SpeedLimit is used to limit the rate at which data is written to the underlying writer.
type SpeedLimit struct {
// Burst is the number of bytes that can be written to the underlying writer per Frame.
Burst int
// Frame is the duration of the frame in which Burst can be written to the underlying writer.
Frame time.Duration
// CheckEveryBytes is the number of bytes written before checking if the speed limit has been exceeded.
CheckEveryBytes int
// Delay is the duration to delay writing if the speed limit has been exceeded during a Write call. (blocking)
Delay time.Duration
}
const fallbackDelay = 100
func regulateSpeedLimit(speedLimit *SpeedLimit) (*SpeedLimit, error) {
if speedLimit.Burst <= 0 || speedLimit.Frame <= 0 {
return nil, errors.New("invalid speed limit")
}
if speedLimit.CheckEveryBytes <= 0 {
speedLimit.CheckEveryBytes = speedLimit.Burst
}
if speedLimit.Delay <= 0 {
speedLimit.Delay = fallbackDelay * time.Millisecond
}
return speedLimit, nil
}
func newSpeedometer(w io.Writer, speedLimit *SpeedLimit, ceiling int64) (*Speedometer, error) {
if w == nil {
return nil, errors.New("writer cannot be nil")
}
var err error
if speedLimit != nil {
if speedLimit, err = regulateSpeedLimit(speedLimit); err != nil {
return nil, err
}
}
return &Speedometer{
w: w,
ceiling: ceiling,
speedLimit: speedLimit,
internal: newAtomics(),
}, nil
}
// NewSpeedometer creates a new Speedometer that wraps the given io.Writer.
// It will not limit the rate at which data is written to the underlying writer, it only measures it.
func NewSpeedometer(w io.Writer) (*Speedometer, error) {
return newSpeedometer(w, nil, -1)
}
// NewLimitedSpeedometer creates a new Speedometer that wraps the given io.Writer.
// If the speed limit is exceeded, writes to the underlying writer will be limited.
// See SpeedLimit for more information.
func NewLimitedSpeedometer(w io.Writer, speedLimit *SpeedLimit) (*Speedometer, error) {
return newSpeedometer(w, speedLimit, -1)
}
// NewCappedSpeedometer creates a new Speedometer that wraps the given io.Writer.
// If len(written) bytes exceeds cap, writes to the underlying writer will be ceased permanently for the Speedometer.
func NewCappedSpeedometer(w io.Writer, capacity int64) (*Speedometer, error) {
return newSpeedometer(w, nil, capacity)
}
// NewCappedLimitedSpeedometer creates a new Speedometer that wraps the given io.Writer.
// It is a combination of NewLimitedSpeedometer and NewCappedSpeedometer.
func NewCappedLimitedSpeedometer(w io.Writer, speedLimit *SpeedLimit, capacity int64) (*Speedometer, error) {
return newSpeedometer(w, speedLimit, capacity)
}
func (s *Speedometer) increment(inc int64) (int, error) {
if s.internal.closed.Load() {
return 0, io.ErrClosedPipe
}
var err error
if s.ceiling > 0 && s.Total()+inc > s.ceiling {
_ = s.Close()
err = ErrLimitReached
inc = s.ceiling - s.Total()
}
s.internal.count.Add(inc)
return int(inc), err
}
// Running returns true if the Speedometer is still running.
func (s *Speedometer) Running() bool {
return !s.internal.closed.Load()
}
// Total returns the total number of bytes written to the underlying writer.
func (s *Speedometer) Total() int64 {
return s.internal.count.Load()
}
// Close stops the Speedometer. No additional writes will be accepted.
func (s *Speedometer) Close() error {
if s.internal.closed.Load() {
return io.ErrClosedPipe
}
s.internal.stop.Do(func() {
s.internal.closed.Store(true)
stopped := time.Now()
birth := s.internal.birth.Load()
duration := stopped.Sub(*birth)
s.internal.duration.Store(&duration)
})
return nil
}
/*func (s *Speedometer) IsSlow() bool {
return s.internal.slow.Load()
}*/
// Rate returns the rate at which data is being written to the underlying writer per second.
func (s *Speedometer) Rate() float64 {
if s.internal.closed.Load() {
return float64(s.Total()) / s.internal.duration.Load().Seconds()
}
return float64(s.Total()) / time.Since(*s.internal.birth.Load()).Seconds()
}
func (s *Speedometer) slowDown() error {
switch {
case s.speedLimit == nil:
return nil
case s.speedLimit.Burst <= 0 || s.speedLimit.Frame <= 0,
s.speedLimit.CheckEveryBytes <= 0, s.speedLimit.Delay <= 0:
return errors.New("invalid speed limit")
default:
//
}
if s.Total()%int64(s.speedLimit.CheckEveryBytes) != 0 {
return nil
}
s.internal.slow.Store(true)
for s.Rate() > float64(s.speedLimit.Burst)/s.speedLimit.Frame.Seconds() {
time.Sleep(s.speedLimit.Delay)
}
s.internal.slow.Store(false)
return nil
}
// Write writes p to the underlying writer, following all defined speed limits.
func (s *Speedometer) Write(p []byte) (n int, err error) {
if s.internal.closed.Load() {
return 0, io.ErrClosedPipe
}
s.internal.start.Do(func() {
now := time.Now()
s.internal.birth.Store(&now)
})
// if no speed limit, just write and record
if s.speedLimit == nil {
n, err = s.w.Write(p)
if err != nil {
return n, fmt.Errorf("error writing to underlying writer: %w", err)
}
return s.increment(int64(len(p)))
}
var (
wErr error
accepted int
)
accepted, wErr = s.increment(int64(len(p)))
if wErr != nil {
return 0, fmt.Errorf("error incrementing: %w", wErr)
}
if sErr := s.slowDown(); sErr != nil {
return 0, fmt.Errorf("error slowing down: %w", sErr)
}
var iErr error
if n, iErr = s.w.Write(p[:accepted]); iErr != nil {
return n, fmt.Errorf("error writing to underlying writer: %w", iErr)
}
return
}

@ -0,0 +1,393 @@
package util
import (
"bytes"
"errors"
"fmt"
"io"
"net"
"sync"
"sync/atomic"
"testing"
"time"
)
type testWriter struct {
t *testing.T
total int64
}
func (w *testWriter) Write(p []byte) (n int, err error) {
atomic.AddInt64(&w.total, int64(len(p)))
return len(p), nil
}
func writeStuff(t *testing.T, target io.Writer, count int) error {
t.Helper()
write := func() error {
_, err := target.Write([]byte("a"))
if err != nil {
return fmt.Errorf("error writing: %w", err)
}
return nil
}
if count < 0 {
var err error
for err = write(); err == nil; err = write() {
time.Sleep(5 * time.Millisecond)
}
return err
}
for i := 0; i < count; i++ {
if err := write(); err != nil {
return err
}
}
return nil
}
//nolint:funlen
func Test_Speedometer(t *testing.T) {
type results struct {
total int64
written int
rate float64
err error
}
isIt := func(want, have results) {
t.Helper()
if have.total != want.total {
t.Errorf("total: want %d, have %d", want.total, have.total)
}
if have.written != want.written {
t.Errorf("written: want %d, have %d", want.written, have.written)
}
if have.rate != want.rate {
t.Errorf("rate: want %f, have %f", want.rate, have.rate)
}
if !errors.Is(have.err, want.err) {
t.Errorf("wantErr: want %v, have %v", want.err, have.err)
}
}
var (
errChan = make(chan error, 10)
)
t.Run("EarlyClose", func(t *testing.T) {
var (
err error
cnt int
)
t.Parallel()
sp, nerr := NewSpeedometer(&testWriter{t: t})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
go func() {
errChan <- writeStuff(t, sp, -1)
}()
time.Sleep(1 * time.Second)
if closeErr := sp.Close(); closeErr != nil {
t.Errorf("wantErr: want %v, have %v", nil, closeErr)
}
err = <-errChan
if !errors.Is(err, io.ErrClosedPipe) {
t.Errorf("wantErr: want %v, have %v", io.ErrClosedPipe, err)
}
cnt, err = sp.Write([]byte("a"))
isIt(results{err: io.ErrClosedPipe, written: 0}, results{err: err, written: cnt})
})
t.Run("Basic", func(t *testing.T) {
var (
err error
cnt int
)
t.Parallel()
sp, nerr := NewSpeedometer(&testWriter{t: t})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
cnt, err = sp.Write([]byte("a"))
isIt(results{err: nil, written: 1, total: 1}, results{err: err, written: cnt, total: sp.Total()})
cnt, err = sp.Write([]byte("aa"))
isIt(results{err: nil, written: 2, total: 3}, results{err: err, written: cnt, total: sp.Total()})
cnt, err = sp.Write([]byte("a"))
isIt(results{err: nil, written: 1, total: 4}, results{err: err, written: cnt, total: sp.Total()})
cnt, err = sp.Write([]byte("a"))
isIt(results{err: nil, written: 1, total: 5}, results{err: err, written: cnt, total: sp.Total()})
})
t.Run("ConcurrentWrites", func(t *testing.T) {
var (
err error
)
count := int64(0)
sp, nerr := NewSpeedometer(&testWriter{t: t})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
wg := &sync.WaitGroup{}
wg.Add(100)
for i := 0; i < 100; i++ {
go func() {
var counted int
var gerr error
counted, gerr = sp.Write([]byte("a"))
if gerr != nil {
t.Errorf("unexpected error: %v", err)
}
atomic.AddInt64(&count, int64(counted))
wg.Done()
}()
}
wg.Wait()
isIt(results{err: nil, written: 100, total: 100},
results{err: err, written: int(atomic.LoadInt64(&count)), total: sp.Total()})
})
t.Run("GottaGoFast", func(t *testing.T) {
t.Parallel()
var (
err error
)
sp, nerr := NewSpeedometer(&testWriter{t: t})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
go func() {
errChan <- writeStuff(t, sp, -1)
}()
var count = 0
for sp.Running() {
select {
case err = <-errChan:
if !errors.Is(err, io.ErrClosedPipe) {
t.Errorf("unexpected error: %v", err)
} else {
if count < 5 {
t.Errorf("too few iterations: %d", count)
}
t.Logf("final rate: %v per second", sp.Rate())
}
default:
if count > 5 {
_ = sp.Close()
}
time.Sleep(100 * time.Millisecond)
t.Logf("rate: %v per second", sp.Rate())
count++
}
}
})
// test limiter with speedlimit
t.Run("CantGoFast", func(t *testing.T) {
t.Parallel()
t.Run("10BytesASecond", func(t *testing.T) {
t.Parallel()
var (
err error
)
sp, nerr := NewLimitedSpeedometer(&testWriter{t: t}, &SpeedLimit{
Burst: 10,
Frame: time.Second,
CheckEveryBytes: 1,
Delay: 100 * time.Millisecond,
})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
for i := 0; i < 15; i++ {
if _, err = sp.Write([]byte("a")); err != nil {
t.Errorf("unexpected error: %v", err)
}
/*if sp.IsSlow() {
t.Errorf("unexpected slow state")
}*/
t.Logf("rate: %v per second", sp.Rate())
if sp.Rate() > 10 {
t.Errorf("speeding in a school zone (expected under %d): %v", sp.speedLimit.Burst, sp.Rate())
}
}
})
t.Run("1000BytesPer5SecondsMeasuredEvery5000Bytes", func(t *testing.T) {
t.Parallel()
var (
err error
)
sp, nerr := NewLimitedSpeedometer(&testWriter{t: t}, &SpeedLimit{
Burst: 1000,
Frame: 2 * time.Second,
CheckEveryBytes: 5000,
Delay: 500 * time.Millisecond,
})
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
for i := 0; i < 4999; i++ {
if _, err = sp.Write([]byte("a")); err != nil {
t.Errorf("unexpected error: %v", err)
}
if i%1000 == 0 {
t.Logf("rate: %v per second", sp.Rate())
}
if sp.Rate() < 1000 {
t.Errorf("shouldn't have slowed down yet (expected over %d): %v", sp.speedLimit.Burst, sp.Rate())
}
}
if _, err = sp.Write([]byte("a")); err != nil {
t.Errorf("unexpected error: %v", err)
}
for i := 0; i < 10; i++ {
if _, err = sp.Write([]byte("a")); err != nil {
t.Errorf("unexpected error: %v", err)
}
t.Logf("rate: %v per second", sp.Rate())
if sp.Rate() > 1000 {
t.Errorf("speeding in a school zone (expected under %d): %v", sp.speedLimit.Burst, sp.Rate())
}
}
})
})
// test capped speedometer
t.Run("OnlyALittle", func(t *testing.T) {
t.Parallel()
var (
err error
)
sp, nerr := NewCappedSpeedometer(&testWriter{t: t}, 1024)
if nerr != nil {
t.Errorf("unexpected error: %v", nerr)
}
for i := 0; i < 1024; i++ {
if _, err = sp.Write([]byte("a")); err != nil {
t.Errorf("unexpected error: %v", err)
}
if sp.Total() > 1024 {
t.Errorf("shouldn't have written more than 1024 bytes")
}
}
if _, err = sp.Write([]byte("a")); err == nil {
t.Errorf("expected error when writing over capacity")
}
})
t.Run("SynSynAckAck", func(t *testing.T) {
t.Parallel()
var (
server net.Listener
err error
)
//goland:noinspection GoCommentLeadingSpace
if server, err = net.Listen("tcp", ":8080"); err != nil { // #nosec:G102 - this is a unit test.
t.Fatalf("Failed to start server: %v", err)
}
defer func(server net.Listener) {
if cErr := server.Close(); cErr != nil {
t.Errorf("Failed to close server: %v", err)
}
}(server)
go func() {
var (
conn net.Conn
aErr error
)
if conn, aErr = server.Accept(); aErr != nil {
t.Errorf("Failed to accept connection: %v", err)
}
t.Logf("Accepted connection from %s", conn.RemoteAddr().String())
defer func(conn net.Conn) {
if cErr := conn.Close(); cErr != nil {
t.Errorf("Failed to close connection: %v", err)
}
}(conn)
speedLimit := &SpeedLimit{
Burst: 512,
Frame: time.Second,
CheckEveryBytes: 1,
Delay: 10 * time.Millisecond,
}
var (
speedometer *Speedometer
sErr error
)
if speedometer, sErr = NewCappedLimitedSpeedometer(conn, speedLimit, 4096); sErr != nil {
t.Errorf("Failed to create speedometer: %v", sErr)
}
buf := make([]byte, 1024)
for i := range buf {
targ := byte('E')
if i%2 == 0 {
targ = byte('e')
}
buf[i] = targ
}
for {
n, wErr := speedometer.Write(buf)
switch {
case errors.Is(wErr, io.EOF), errors.Is(wErr, ErrLimitReached):
return
case wErr != nil:
t.Errorf("Failed to write: %v", wErr)
case n != len(buf):
t.Errorf("Failed to write all bytes: %d", n)
default:
t.Logf("Wrote %d bytes", n)
}
}
}()
var (
client net.Conn
aErr error
)
if client, aErr = net.Dial("tcp", "localhost:8080"); aErr != nil {
t.Fatalf("Failed to connect to server: %v", err)
}
defer func(client net.Conn) {
if clErr := client.Close(); clErr != nil {
t.Errorf("Failed to close client: %v", err)
}
}(client)
buf := &bytes.Buffer{}
startTime := time.Now()
n, cpErr := io.Copy(buf, client)
if cpErr != nil {
t.Errorf("Failed to copy: %v", cpErr)
}
duration := time.Since(startTime)
if buf.Len() == 0 || n == 0 {
t.Fatalf("No data received")
}
rate := measureRate(t, n, duration)
if rate > 512.0 {
t.Fatalf("Rate exceeded: got %f, expected <= 100.0", rate)
}
})
}
func measureRate(t *testing.T, received int64, duration time.Duration) float64 {
t.Helper()
return float64(received) / duration.Seconds()
}