8 Commits

Author SHA1 Message Date
b612 3add9183b3 fix:修复一些边界性问题 2026-03-21 19:15:43 +08:00
b612 3c83d1d39f bug fix:修正了InputMsg在空字符串场景下的问题 2024-08-30 14:51:51 +08:00
b612 d528a583f4 update go mod 2024-08-18 17:18:10 +08:00
b612 a986fafa09 move starqueue from starnet to stario 2024-08-18 17:17:20 +08:00
b612 8e09b7eec0 update go.mod 2024-03-10 14:01:31 +08:00
b612 609bdf7079 add signal control for password input 2023-02-11 16:55:57 +08:00
b612 95503ec65a bug fix 2022-03-18 14:51:57 +08:00
b612 538445a201 add InputMsg slice support 2022-03-18 13:49:38 +08:00
13 changed files with 1412 additions and 233 deletions
+7
View File
@@ -0,0 +1,7 @@
.sentrux/
agent_readme.md
target.md
.gocache
.gocache/
.tmp_*/
.idea
+137 -94
View File
@@ -2,148 +2,191 @@ package stario
import (
"errors"
"fmt"
"io"
"os"
"runtime"
"sync"
"sync/atomic"
"time"
)
var ErrStarBufferInvalidCapacity = errors.New("star buffer capacity must be greater than zero")
var ErrStarBufferClosed = errors.New("star buffer closed")
var ErrStarBufferWriteClosed = errors.New("star buffer write closed")
type StarBuffer struct {
io.Reader
io.Writer
io.Closer
datas []byte
pStart uint64
pEnd uint64
cap uint64
isClose atomic.Value
isEnd atomic.Value
rmu sync.Mutex
wmu sync.Mutex
datas []byte
pStart uint64
pEnd uint64
size uint64
cap uint64
isClose bool
isWriteEnd bool
mu sync.Mutex
notEmpty *sync.Cond
notFull *sync.Cond
}
func NewStarBuffer(cap uint64) *StarBuffer {
rtnBuffer := new(StarBuffer)
rtnBuffer.cap = cap
rtnBuffer.datas = make([]byte, cap)
rtnBuffer.isClose.Store(false)
rtnBuffer.isEnd.Store(false)
return rtnBuffer
func NewStarBuffer(cap uint64) (*StarBuffer, error) {
if cap == 0 {
return nil, ErrStarBufferInvalidCapacity
}
rtnBuffer := &StarBuffer{
cap: cap,
datas: make([]byte, cap),
}
rtnBuffer.notEmpty = sync.NewCond(&rtnBuffer.mu)
rtnBuffer.notFull = sync.NewCond(&rtnBuffer.mu)
return rtnBuffer, nil
}
func (star *StarBuffer) Free() uint64 {
return star.cap - star.Len()
star.mu.Lock()
defer star.mu.Unlock()
return star.cap - star.size
}
func (star *StarBuffer) Cap() uint64 {
star.mu.Lock()
defer star.mu.Unlock()
return star.cap
}
func (star *StarBuffer) Len() uint64 {
if star.pEnd >= star.pStart {
return star.pEnd - star.pStart
}
return star.pEnd - star.pStart + star.cap
star.mu.Lock()
defer star.mu.Unlock()
return star.size
}
func (star *StarBuffer) getByte() (byte, error) {
if star.isClose.Load().(bool) || (star.Len() == 0 && star.isEnd.Load().(bool)) {
star.mu.Lock()
defer star.mu.Unlock()
for star.size == 0 && !star.isWriteEnd && !star.isClose {
star.notEmpty.Wait()
}
if star.size == 0 {
return 0, io.EOF
}
if star.Len() == 0 {
return 0, os.ErrNotExist
}
nowPtr := star.pStart
nextPtr := star.pStart + 1
if nextPtr >= star.cap {
nextPtr = 0
}
data := star.datas[nowPtr]
ok := atomic.CompareAndSwapUint64(&star.pStart, nowPtr, nextPtr)
if !ok {
return 0, os.ErrInvalid
}
data := star.datas[star.pStart]
star.pStart = (star.pStart + 1) % star.cap
star.size--
star.notFull.Broadcast()
return data, nil
}
func (star *StarBuffer) putByte(data byte) error {
if star.isClose.Load().(bool) {
return io.EOF
star.mu.Lock()
defer star.mu.Unlock()
for star.size == star.cap && !star.isClose && !star.isWriteEnd {
star.notFull.Wait()
}
nowPtr := star.pEnd
kariEnd := nowPtr + 1
if kariEnd == star.cap {
kariEnd = 0
if star.isClose {
return ErrStarBufferClosed
}
if kariEnd == atomic.LoadUint64(&star.pStart) {
for {
time.Sleep(time.Microsecond)
runtime.Gosched()
if kariEnd != atomic.LoadUint64(&star.pStart) {
break
}
}
if star.isWriteEnd {
return ErrStarBufferWriteClosed
}
star.datas[nowPtr] = data
if ok := atomic.CompareAndSwapUint64(&star.pEnd, nowPtr, kariEnd); !ok {
return os.ErrInvalid
star.datas[star.pEnd] = data
star.pEnd = (star.pEnd + 1) % star.cap
star.size++
star.notEmpty.Broadcast()
return nil
}
func (star *StarBuffer) EndWrite() error {
star.mu.Lock()
defer star.mu.Unlock()
if star.isClose {
return ErrStarBufferClosed
}
star.isWriteEnd = true
star.notEmpty.Broadcast()
star.notFull.Broadcast()
return nil
}
func (star *StarBuffer) Close() error {
star.isClose.Store(true)
star.mu.Lock()
defer star.mu.Unlock()
star.isClose = true
star.isWriteEnd = true
star.notEmpty.Broadcast()
star.notFull.Broadcast()
return nil
}
func (star *StarBuffer) Read(buf []byte) (int, error) {
if star.isClose.Load().(bool) || (star.Len() == 0 && star.isEnd.Load().(bool)) {
return 0, io.EOF
}
if buf == nil {
return 0, errors.New("buffer is nil")
}
star.rmu.Lock()
defer star.rmu.Unlock()
var sum int = 0
for i := 0; i < len(buf); i++ {
data, err := star.getByte()
if err != nil {
if err == io.EOF {
return sum, err
}
if err == os.ErrNotExist {
i--
continue
}
return sum, nil
}
buf[i] = data
sum++
if len(buf) == 0 {
return 0, nil
}
star.mu.Lock()
defer star.mu.Unlock()
for star.size == 0 && !star.isWriteEnd && !star.isClose {
star.notEmpty.Wait()
}
if star.size == 0 {
return 0, io.EOF
}
sum := minInt(len(buf), int(star.size))
first := minInt(sum, int(star.cap-star.pStart))
copy(buf, star.datas[star.pStart:star.pStart+uint64(first)])
second := sum - first
if second > 0 {
copy(buf[first:], star.datas[:second])
}
star.pStart = (star.pStart + uint64(sum)) % star.cap
star.size -= uint64(sum)
star.notFull.Broadcast()
return sum, nil
}
func (star *StarBuffer) Write(bts []byte) (int, error) {
if bts == nil && !star.isEnd.Load().(bool) {
star.isEnd.Store(true)
if bts == nil {
return 0, star.EndWrite()
}
if len(bts) == 0 {
return 0, nil
}
if bts == nil || star.isClose.Load().(bool) {
return 0, io.EOF
}
star.wmu.Lock()
defer star.wmu.Unlock()
var sum = 0
for i := 0; i < len(bts); i++ {
err := star.putByte(bts[i])
if err != nil {
fmt.Println("Write bts err:", err)
return sum, err
star.mu.Lock()
defer star.mu.Unlock()
sum := 0
for sum < len(bts) {
for star.size == star.cap && !star.isClose && !star.isWriteEnd {
star.notFull.Wait()
}
sum++
if star.isClose {
if sum == 0 {
return 0, ErrStarBufferClosed
}
return sum, ErrStarBufferClosed
}
if star.isWriteEnd {
if sum == 0 {
return 0, ErrStarBufferWriteClosed
}
return sum, ErrStarBufferWriteClosed
}
space := int(star.cap - star.size)
if space == 0 {
continue
}
n := minInt(len(bts)-sum, space)
first := minInt(n, int(star.cap-star.pEnd))
copy(star.datas[star.pEnd:star.pEnd+uint64(first)], bts[sum:sum+first])
second := n - first
if second > 0 {
copy(star.datas[:second], bts[sum+first:sum+n])
}
star.pEnd = (star.pEnd + uint64(n)) % star.cap
star.size += uint64(n)
sum += n
star.notEmpty.Broadcast()
}
return sum, nil
}
func minInt(a int, b int) int {
if a < b {
return a
}
return b
}
+86 -3
View File
@@ -1,14 +1,91 @@
package stario
import (
"bytes"
"fmt"
"io"
"sync/atomic"
"testing"
"time"
)
func TestNewStarBufferRejectsZeroCapacity(t *testing.T) {
buf, err := NewStarBuffer(0)
if err != ErrStarBufferInvalidCapacity {
t.Fatalf("unexpected error: %v", err)
}
if buf != nil {
t.Fatal("expected nil buffer when capacity is invalid")
}
}
func TestStarBufferEndWriteDrainsThenEOF(t *testing.T) {
buf, err := NewStarBuffer(4)
if err != nil {
t.Fatal(err)
}
if _, err := buf.Write([]byte("abcd")); err != nil {
t.Fatal(err)
}
if err := buf.EndWrite(); err != nil {
t.Fatal(err)
}
got := make([]byte, 4)
n, err := buf.Read(got)
if err != nil {
t.Fatal(err)
}
if n != 4 || !bytes.Equal(got[:n], []byte("abcd")) {
t.Fatalf("unexpected payload: n=%d data=%q", n, got[:n])
}
n, err = buf.Read(got)
if n != 0 || err != io.EOF {
t.Fatalf("expected EOF after draining buffer, got n=%d err=%v", n, err)
}
if _, err := buf.Write([]byte("x")); err != ErrStarBufferWriteClosed {
t.Fatalf("unexpected write error after EndWrite: %v", err)
}
}
func TestStarBufferCloseAllowsDrain(t *testing.T) {
buf, err := NewStarBuffer(4)
if err != nil {
t.Fatal(err)
}
if _, err := buf.Write([]byte("ab")); err != nil {
t.Fatal(err)
}
if err := buf.Close(); err != nil {
t.Fatal(err)
}
got := make([]byte, 2)
n, err := buf.Read(got)
if err != nil {
t.Fatal(err)
}
if n != 2 || !bytes.Equal(got[:n], []byte("ab")) {
t.Fatalf("unexpected payload after close: n=%d data=%q", n, got[:n])
}
n, err = buf.Read(got)
if n != 0 || err != io.EOF {
t.Fatalf("expected EOF after draining closed buffer, got n=%d err=%v", n, err)
}
if _, err := buf.Write([]byte("x")); err != ErrStarBufferClosed {
t.Fatalf("unexpected write error after Close: %v", err)
}
}
func Test_Circle(t *testing.T) {
buf := NewStarBuffer(2048)
buf, err := NewStarBuffer(2048)
if err != nil {
t.Fatal(err)
}
go func() {
for {
//fmt.Println("write start")
@@ -37,7 +114,10 @@ func Test_Circle(t *testing.T) {
}
func Test_Circle_Speed(t *testing.T) {
buf := NewStarBuffer(1048976)
buf, err := NewStarBuffer(1048976)
if err != nil {
t.Fatal(err)
}
count := uint64(0)
for i := 1; i <= 10; i++ {
go func() {
@@ -61,7 +141,10 @@ func Test_Circle_Speed(t *testing.T) {
}
func Test_Circle_Speed2(t *testing.T) {
buf := NewStarBuffer(8192)
buf, err := NewStarBuffer(8192)
if err != nil {
t.Fatal(err)
}
count := uint64(0)
for i := 1; i <= 10; i++ {
go func() {
+1 -1
View File
@@ -2,4 +2,4 @@ module b612.me/stario
go 1.16
require golang.org/x/crypto v0.0.0-20220313003712-b769efc7c000
require golang.org/x/crypto v0.26.0
+63 -7
View File
@@ -1,11 +1,67 @@
golang.org/x/crypto v0.0.0-20220313003712-b769efc7c000 h1:SL+8VVnkqyshUSz5iNnXtrBQzvFF2SkROm6t5RczFAE=
golang.org/x/crypto v0.0.0-20220313003712-b769efc7c000/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw=
golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 h1:SrN+KX8Art/Sf4HNj6Zcz06G7VEz+7w9tdXTPOZ7+l4=
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.23.0 h1:YfKFowiIMvtgl1UERQoTPPToxltDeZfbj4H7dVUCwmM=
golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU=
golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+331 -98
View File
@@ -3,126 +3,215 @@ package stario
import (
"bufio"
"fmt"
"golang.org/x/crypto/ssh/terminal"
"os"
"runtime"
"strconv"
"strings"
"golang.org/x/crypto/ssh/terminal"
)
type InputMsg struct {
msg string
err error
msg string
err error
skipSliceSigErr bool
}
type rawTerminalSession struct {
fd int
state *terminal.State
reader *bufio.Reader
redrawHint string
printNewline bool
}
func Passwd(hint string, defaultVal string) InputMsg {
return passwd(hint, defaultVal, "")
return passwd(hint, defaultVal, "", true)
}
func PasswdWithMask(hint string, defaultVal string, mask string) InputMsg {
return passwd(hint, defaultVal, mask)
return passwd(hint, defaultVal, mask, true)
}
func PasswdResponseSignal(hint string, defaultVal string) InputMsg {
return passwd(hint, defaultVal, "", true)
}
func PasswdResponseSignalWithMask(hint string, defaultVal string, mask string) InputMsg {
return passwd(hint, defaultVal, mask, true)
}
func MessageBoxRaw(hint string, defaultVal string) InputMsg {
return messageBox(hint, defaultVal)
}
func messageBox(hint string, defaultVal string) InputMsg {
var ioBuf []rune
func newRawTerminalSession(hint string, printNewline bool) (*rawTerminalSession, error) {
if hint != "" {
fmt.Print(hint)
}
if strings.Index(hint, "\n") >= 0 {
hint = strings.TrimSpace(hint[strings.LastIndex(hint, "\n"):])
}
fd := int(os.Stdin.Fd())
state, err := terminal.MakeRaw(fd)
if err != nil {
return InputMsg{"", err}
return nil, err
}
defer fmt.Println()
defer terminal.Restore(fd, state)
inputReader := bufio.NewReader(os.Stdin)
for {
b, _, err := inputReader.ReadRune()
if err != nil {
return InputMsg{"", err}
}
if b == 0x0d {
strValue := strings.TrimSpace(string(ioBuf))
if len(strValue) == 0 {
strValue = defaultVal
}
return InputMsg{strValue, nil}
}
if b == 0x08 || b == 0x7F {
if len(ioBuf) > 0 {
ioBuf = ioBuf[:len(ioBuf)-1]
}
fmt.Print("\r")
for i := 0; i < len(ioBuf)+2+len(hint); i++ {
fmt.Print(" ")
}
} else {
ioBuf = append(ioBuf, b)
}
fmt.Print("\r")
if hint != "" {
fmt.Print(hint)
}
fmt.Print(string(ioBuf))
return &rawTerminalSession{
fd: fd,
state: state,
reader: bufio.NewReader(os.Stdin),
redrawHint: promptRedrawHint(hint),
printNewline: printNewline,
}, nil
}
func (session *rawTerminalSession) Close() {
if session == nil || session.state == nil {
return
}
_ = terminal.Restore(session.fd, session.state)
if session.printNewline {
fmt.Println()
}
}
func passwd(hint string, defaultVal string, mask string) InputMsg {
var ioBuf []rune
func (session *rawTerminalSession) Restore() error {
if session == nil || session.state == nil {
return nil
}
return terminal.Restore(session.fd, session.state)
}
func promptRedrawHint(hint string) string {
if strings.Index(hint, "\n") >= 0 {
hint = hint[strings.LastIndex(hint, "\n")+1:]
}
return strings.TrimSpace(hint)
}
func finalizeInputValue(raw string, defaultVal string) string {
raw = strings.TrimSpace(raw)
if len(raw) == 0 {
return defaultVal
}
return raw
}
func renderRawEcho(ioBuf []rune, mask string) string {
if mask == "" {
return string(ioBuf)
}
return strings.Repeat(mask, len(ioBuf))
}
func redrawPromptLine(hint string, echo string, lastWidth int) int {
nowWidth := stringDisplayWidth(hint) + stringDisplayWidth(echo)
clearWidth := lastWidth
if nowWidth > clearWidth {
clearWidth = nowWidth
}
fmt.Print("\r")
if clearWidth > 0 {
fmt.Print(strings.Repeat(" ", clearWidth))
fmt.Print("\r")
}
if hint != "" {
fmt.Print(hint)
}
if strings.Index(hint, "\n") >= 0 {
hint = strings.TrimSpace(hint[strings.LastIndex(hint, "\n"):])
if echo != "" {
fmt.Print(echo)
}
fd := int(os.Stdin.Fd())
state, err := terminal.MakeRaw(fd)
return nowWidth
}
func stringDisplayWidth(text string) int {
width := 0
for _, r := range text {
width += runeDisplayWidth(r)
}
return width
}
func runeDisplayWidth(r rune) int {
switch {
case r == '\t':
return 4
case r < 0x20 || (r >= 0x7f && r < 0xa0):
return 0
case isWideRune(r):
return 2
default:
return 1
}
}
func isWideRune(r rune) bool {
return r >= 0x1100 && (r <= 0x115f ||
r == 0x2329 || r == 0x232a ||
(r >= 0x2e80 && r <= 0xa4cf && r != 0x303f) ||
(r >= 0xac00 && r <= 0xd7a3) ||
(r >= 0xf900 && r <= 0xfaff) ||
(r >= 0xfe10 && r <= 0xfe19) ||
(r >= 0xfe30 && r <= 0xfe6f) ||
(r >= 0xff00 && r <= 0xff60) ||
(r >= 0xffe0 && r <= 0xffe6) ||
(r >= 0x1f300 && r <= 0x1f64f) ||
(r >= 0x1f900 && r <= 0x1f9ff) ||
(r >= 0x20000 && r <= 0x3fffd))
}
func rawLineInput(hint string, defaultVal string, mask string, handleSignal bool) InputMsg {
session, err := newRawTerminalSession(hint, true)
if err != nil {
return InputMsg{"", err}
return InputMsg{msg: "", err: err}
}
defer fmt.Println()
defer terminal.Restore(fd, state)
inputReader := bufio.NewReader(os.Stdin)
defer session.Close()
ioBuf := make([]rune, 0, 16)
lastWidth := 0
for {
b, _, err := inputReader.ReadRune()
b, _, err := session.reader.ReadRune()
if err != nil {
return InputMsg{"", err}
return InputMsg{msg: "", err: err}
}
if b == 0x0d {
strValue := strings.TrimSpace(string(ioBuf))
if len(strValue) == 0 {
strValue = defaultVal
if handleSignal && isSignal(b) {
if runtime.GOOS != "windows" {
if err := session.Restore(); err != nil {
return InputMsg{msg: "", err: err}
}
}
return InputMsg{strValue, nil}
if err := signal(b); err != nil {
return InputMsg{msg: "", err: err}
}
continue
}
if b == 0x08 || b == 0x7F {
switch b {
case 0x0d, 0x0a:
return InputMsg{msg: finalizeInputValue(string(ioBuf), defaultVal), err: nil}
case 0x08, 0x7F:
if len(ioBuf) > 0 {
ioBuf = ioBuf[:len(ioBuf)-1]
}
fmt.Print("\r")
for i := 0; i < len(ioBuf)+2+len(hint); i++ {
fmt.Print(" ")
}
} else {
default:
ioBuf = append(ioBuf, b)
}
fmt.Print("\r")
if hint != "" {
fmt.Print(hint)
}
for i := 0; i < len(ioBuf); i++ {
fmt.Print(mask)
}
lastWidth = redrawPromptLine(session.redrawHint, renderRawEcho(ioBuf, mask), lastWidth)
}
}
func messageBox(hint string, defaultVal string) InputMsg {
return rawLineInput(hint, defaultVal, "", false)
}
func isSignal(s rune) bool {
switch s {
case 0x03, 0x1a, 0x1c:
return true
default:
return false
}
}
func passwd(hint string, defaultVal string, mask string, handleSignal bool) InputMsg {
return rawLineInput(hint, defaultVal, mask, handleSignal)
}
func MessageBox(hint string, defaultVal string) InputMsg {
if hint != "" {
fmt.Print(hint)
@@ -130,13 +219,15 @@ func MessageBox(hint string, defaultVal string) InputMsg {
inputReader := bufio.NewReader(os.Stdin)
str, err := inputReader.ReadString('\n')
if err != nil {
return InputMsg{"", err}
return InputMsg{msg: str, err: err}
}
str = strings.TrimSpace(str)
if len(str) == 0 {
str = defaultVal
}
return InputMsg{str, nil}
str = finalizeInputValue(str, defaultVal)
return InputMsg{msg: str, err: err}
}
func (im *InputMsg) IgnoreSliceParseError(i bool) *InputMsg {
im.skipSliceSigErr = i
return im
}
func (im InputMsg) String() (string, error) {
@@ -151,6 +242,38 @@ func (im InputMsg) MustString() string {
return res
}
func (im InputMsg) SliceString(sep string) ([]string, error) {
if im.err != nil {
return nil, im.err
}
if len(strings.TrimSpace(im.msg)) == 0 {
return []string{}, nil
}
return strings.Split(im.msg, sep), nil
}
func (im InputMsg) MustSliceString(sep string) []string {
res, _ := im.SliceString(sep)
return res
}
func (im InputMsg) sliceFn(sep string, fn func(string) (interface{}, error)) ([]interface{}, error) {
var res []interface{}
data, err := im.SliceString(sep)
if err != nil {
return res, err
}
for _, v := range data {
code, err := fn(v)
if err != nil && !im.skipSliceSigErr {
return nil, err
} else if err == nil {
res = append(res, code)
}
}
return res, nil
}
func (im InputMsg) Int() (int, error) {
if im.err != nil {
return 0, im.err
@@ -158,6 +281,22 @@ func (im InputMsg) Int() (int, error) {
return strconv.Atoi(im.msg)
}
func (im InputMsg) SliceInt(sep string) ([]int, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.Atoi(v)
})
var res []int
for _, v := range data {
res = append(res, v.(int))
}
return res, err
}
func (im InputMsg) MustSliceInt(sep string) []int {
res, _ := im.SliceInt(sep)
return res
}
func (im InputMsg) MustInt() int {
res, _ := im.Int()
return res
@@ -175,6 +314,22 @@ func (im InputMsg) MustInt64() int64 {
return res
}
func (im InputMsg) SliceInt64(sep string) ([]int64, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.ParseInt(v, 10, 64)
})
var res []int64
for _, v := range data {
res = append(res, v.(int64))
}
return res, err
}
func (im InputMsg) MustSliceInt64(sep string) []int64 {
res, _ := im.SliceInt64(sep)
return res
}
func (im InputMsg) Uint64() (uint64, error) {
if im.err != nil {
return 0, im.err
@@ -186,6 +341,21 @@ func (im InputMsg) MustUint64() uint64 {
res, _ := im.Uint64()
return res
}
func (im InputMsg) SliceUint64(sep string) ([]uint64, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.ParseUint(v, 10, 64)
})
var res []uint64
for _, v := range data {
res = append(res, v.(uint64))
}
return res, err
}
func (im InputMsg) MustSliceUint64(sep string) []uint64 {
res, _ := im.SliceUint64(sep)
return res
}
func (im InputMsg) Bool() (bool, error) {
if im.err != nil {
@@ -199,6 +369,22 @@ func (im InputMsg) MustBool() bool {
return res
}
func (im InputMsg) SliceBool(sep string) ([]bool, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.ParseBool(v)
})
var res []bool
for _, v := range data {
res = append(res, v.(bool))
}
return res, err
}
func (im InputMsg) MustSliceBool(sep string) []bool {
res, _ := im.SliceBool(sep)
return res
}
func (im InputMsg) Float64() (float64, error) {
if im.err != nil {
return 0, im.err
@@ -211,6 +397,22 @@ func (im InputMsg) MustFloat64() float64 {
return res
}
func (im InputMsg) SliceFloat64(sep string) ([]float64, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.ParseFloat(v, 64)
})
var res []float64
for _, v := range data {
res = append(res, v.(float64))
}
return res, err
}
func (im InputMsg) MustSliceFloat64(sep string) []float64 {
res, _ := im.SliceFloat64(sep)
return res
}
func (im InputMsg) Float32() (float32, error) {
if im.err != nil {
return 0, im.err
@@ -224,47 +426,78 @@ func (im InputMsg) MustFloat32() float32 {
return res
}
func (im InputMsg) SliceFloat32(sep string) ([]float32, error) {
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
return strconv.ParseFloat(v, 32)
})
var res []float32
for _, v := range data {
res = append(res, v.(float32))
}
return res, err
}
func (im InputMsg) MustSliceFloat32(sep string) []float32 {
res, _ := im.SliceFloat32(sep)
return res
}
func YesNo(hint string, defaults bool) bool {
res, err := YesNoE(hint, defaults)
if err != nil {
return defaults
}
return res
}
func parseYesNoValue(raw string, defaults bool) (bool, bool) {
raw = strings.TrimSpace(strings.ToUpper(raw))
if raw == "" {
return defaults, true
}
switch []rune(raw)[0] {
case 'Y':
return true, true
case 'N':
return false, true
default:
return false, false
}
}
func YesNoE(hint string, defaults bool) (bool, error) {
for {
res := strings.ToUpper(MessageBox(hint, "").MustString())
if res == "" {
return defaults
res, err := MessageBox(hint, "").String()
if err != nil {
return false, err
}
res = res[0:1]
if res == "Y" {
return true
} else if res == "N" {
return false
if answer, ok := parseYesNoValue(res, defaults); ok {
return answer, nil
}
}
}
func StopUntil(hint string, trigger string, repeat bool) error {
pressLen := len([]rune(trigger))
triggerRunes := []rune(trigger)
pressLen := len(triggerRunes)
if trigger == "" {
pressLen = 1
}
fd := int(os.Stdin.Fd())
if hint != "" {
fmt.Print(hint)
}
state, err := terminal.MakeRaw(fd)
session, err := newRawTerminalSession(hint, false)
if err != nil {
return err
}
defer terminal.Restore(fd, state)
inputReader := bufio.NewReader(os.Stdin)
//ioBuf := make([]byte, pressLen)
defer session.Close()
i := 0
for {
b, _, err := inputReader.ReadRune()
b, _, err := session.reader.ReadRune()
if err != nil {
return err
}
if trigger == "" {
break
}
if b == []rune(trigger)[i] {
if b == triggerRunes[i] {
i++
if i == pressLen {
break
+84
View File
@@ -0,0 +1,84 @@
package stario
import (
"fmt"
"testing"
)
func TestPromptRedrawHint(t *testing.T) {
got := promptRedrawHint("头部提示\n 中文确认: ")
if got != "中文确认:" {
t.Fatalf("unexpected redraw hint: got %q", got)
}
}
func TestStringDisplayWidth(t *testing.T) {
got := stringDisplayWidth("中a[]")
if got != 5 {
t.Fatalf("unexpected display width: got %d want 5", got)
}
}
func TestParseYesNoValue(t *testing.T) {
cases := []struct {
name string
input string
defaults bool
want bool
ok bool
}{
{name: "default", input: " ", defaults: true, want: true, ok: true},
{name: "yes", input: "yes", defaults: false, want: true, ok: true},
{name: "no", input: "No", defaults: true, want: false, ok: true},
{name: "invalid", input: "maybe", defaults: false, want: false, ok: false},
}
for _, tc := range cases {
got, ok := parseYesNoValue(tc.input, tc.defaults)
if got != tc.want || ok != tc.ok {
t.Fatalf("%s: got (%v, %v) want (%v, %v)", tc.name, got, ok, tc.want, tc.ok)
}
}
}
func Test_Slice(t *testing.T) {
var data = InputMsg{
msg: "true,false,true,true,false,0,1,hello",
err: nil,
skipSliceSigErr: false,
}
res, err := data.IgnoreSliceParseError(true).SliceBool(",")
if err != nil {
fmt.Println(res)
t.Fatal(err)
}
if len(res) == 0 {
t.Fatal(res)
}
fmt.Println(res)
}
func TestSliceMsg(t *testing.T) {
var data = InputMsg{
msg: "",
err: nil,
skipSliceSigErr: false,
}
res, err := data.SliceString(",")
if err != nil {
fmt.Println(res)
t.Fatal(err)
}
if len(res) != 0 {
t.Fatal(res)
}
fmt.Println(len(res))
res2, err := data.SliceInt64(",")
if err != nil {
fmt.Println(res2)
t.Fatal(err)
}
if len(res2) != 0 {
t.Fatal(res2)
}
fmt.Println(len(res2))
}
+325
View File
@@ -0,0 +1,325 @@
package stario
import (
"bytes"
"context"
"encoding/binary"
"errors"
"fmt"
"os"
"sync"
"time"
)
var ErrDeadlineExceeded error = errors.New("deadline exceeded")
// 识别头
var header = []byte{11, 27, 19, 96, 12, 25, 02, 20}
// MsgQueue 为基本的信息单位
type MsgQueue struct {
ID uint16
Msg []byte
Conn interface{}
}
// StarQueue 为流数据中的消息队列分发
type StarQueue struct {
maxLength uint32
count int64
Encode bool
msgID uint16
msgPool chan MsgQueue
unFinMsg sync.Map
lastID int //= -1
ctx context.Context
cancel context.CancelFunc
duration time.Duration
EncodeFunc func([]byte) []byte
DecodeFunc func([]byte) []byte
//restoreMu sync.Mutex
}
func NewQueueCtx(ctx context.Context, count int64, maxMsgLength uint32) *StarQueue {
var q StarQueue
q.Encode = false
q.count = count
q.maxLength = maxMsgLength
q.msgPool = make(chan MsgQueue, count)
if ctx == nil {
q.ctx, q.cancel = context.WithCancel(context.Background())
} else {
q.ctx, q.cancel = context.WithCancel(ctx)
}
q.duration = 0
return &q
}
func NewQueueWithCount(count int64) *StarQueue {
return NewQueueCtx(nil, count, 0)
}
// NewQueue 建立一个新消息队列
func NewQueue() *StarQueue {
return NewQueueWithCount(32)
}
// Uint32ToByte 4位uint32转byte
func Uint32ToByte(src uint32) []byte {
res := make([]byte, 4)
res[3] = uint8(src)
res[2] = uint8(src >> 8)
res[1] = uint8(src >> 16)
res[0] = uint8(src >> 24)
return res
}
// ByteToUint32 byte转4位uint32
func ByteToUint32(src []byte) uint32 {
var res uint32
buffer := bytes.NewBuffer(src)
binary.Read(buffer, binary.BigEndian, &res)
return res
}
// Uint16ToByte 2位uint16转byte
func Uint16ToByte(src uint16) []byte {
res := make([]byte, 2)
res[1] = uint8(src)
res[0] = uint8(src >> 8)
return res
}
// ByteToUint16 用于byte转uint16
func ByteToUint16(src []byte) uint16 {
var res uint16
buffer := bytes.NewBuffer(src)
binary.Read(buffer, binary.BigEndian, &res)
return res
}
// BuildMessage 生成编码后的信息用于发送
func (q *StarQueue) BuildMessage(src []byte) []byte {
var buff bytes.Buffer
q.msgID++
if q.Encode {
src = q.EncodeFunc(src)
}
length := uint32(len(src))
buff.Write(header)
buff.Write(Uint32ToByte(length))
buff.Write(Uint16ToByte(q.msgID))
buff.Write(src)
return buff.Bytes()
}
// BuildHeader 生成编码后的Header用于发送
func (q *StarQueue) BuildHeader(length uint32) []byte {
var buff bytes.Buffer
q.msgID++
buff.Write(header)
buff.Write(Uint32ToByte(length))
buff.Write(Uint16ToByte(q.msgID))
return buff.Bytes()
}
type unFinMsg struct {
ID uint16
LengthRecv uint32
// HeaderMsg 信息头,应当为14位:8位识别码+4位长度码+2位id
HeaderMsg []byte
RecvMsg []byte
}
func (q *StarQueue) push2list(msg MsgQueue) {
q.msgPool <- msg
}
// ParseMessage 用于解析收到的msg信息
func (q *StarQueue) ParseMessage(msg []byte, conn interface{}) error {
return q.parseMessage(msg, conn)
}
// parseMessage 用于解析收到的msg信息
func (q *StarQueue) parseMessage(msg []byte, conn interface{}) error {
tmp, ok := q.unFinMsg.Load(conn)
if ok { //存在未完成的信息
lastMsg := tmp.(*unFinMsg)
headerLen := len(lastMsg.HeaderMsg)
if headerLen < 14 { //未完成头标题
//传输的数据不能填充header头
if len(msg) < 14-headerLen {
//加入header头并退出
lastMsg.HeaderMsg = bytesMerge(lastMsg.HeaderMsg, msg)
q.unFinMsg.Store(conn, lastMsg)
return nil
}
//获取14字节完整的header
header := msg[0 : 14-headerLen]
lastMsg.HeaderMsg = bytesMerge(lastMsg.HeaderMsg, header)
//检查收到的header是否为认证header
//若不是,丢弃并重新来过
if !checkHeader(lastMsg.HeaderMsg[0:8]) {
q.unFinMsg.Delete(conn)
if len(msg) == 0 {
return nil
}
return q.parseMessage(msg, conn)
}
//获得本数据包长度
lastMsg.LengthRecv = ByteToUint32(lastMsg.HeaderMsg[8:12])
if q.maxLength != 0 && lastMsg.LengthRecv > q.maxLength {
q.unFinMsg.Delete(conn)
return fmt.Errorf("msg length is %d ,too large than %d", lastMsg.LengthRecv, q.maxLength)
}
//获得本数据包ID
lastMsg.ID = ByteToUint16(lastMsg.HeaderMsg[12:14])
//存入列表
q.unFinMsg.Store(conn, lastMsg)
msg = msg[14-headerLen:]
if uint32(len(msg)) < lastMsg.LengthRecv {
lastMsg.RecvMsg = msg
q.unFinMsg.Store(conn, lastMsg)
return nil
}
if uint32(len(msg)) >= lastMsg.LengthRecv {
lastMsg.RecvMsg = msg[0:lastMsg.LengthRecv]
if q.Encode {
lastMsg.RecvMsg = q.DecodeFunc(lastMsg.RecvMsg)
}
msg = msg[lastMsg.LengthRecv:]
storeMsg := MsgQueue{
ID: lastMsg.ID,
Msg: lastMsg.RecvMsg,
Conn: conn,
}
//q.restoreMu.Lock()
q.push2list(storeMsg)
//q.restoreMu.Unlock()
q.unFinMsg.Delete(conn)
return q.parseMessage(msg, conn)
}
} else {
lastID := int(lastMsg.LengthRecv) - len(lastMsg.RecvMsg)
if lastID < 0 {
q.unFinMsg.Delete(conn)
return q.parseMessage(msg, conn)
}
if len(msg) >= lastID {
lastMsg.RecvMsg = bytesMerge(lastMsg.RecvMsg, msg[0:lastID])
if q.Encode {
lastMsg.RecvMsg = q.DecodeFunc(lastMsg.RecvMsg)
}
storeMsg := MsgQueue{
ID: lastMsg.ID,
Msg: lastMsg.RecvMsg,
Conn: conn,
}
q.push2list(storeMsg)
q.unFinMsg.Delete(conn)
if len(msg) == lastID {
return nil
}
msg = msg[lastID:]
return q.parseMessage(msg, conn)
}
lastMsg.RecvMsg = bytesMerge(lastMsg.RecvMsg, msg)
q.unFinMsg.Store(conn, lastMsg)
return nil
}
}
if len(msg) == 0 {
return nil
}
var start int
if start = searchHeader(msg); start == -1 {
return errors.New("data format error")
}
msg = msg[start:]
lastMsg := unFinMsg{}
q.unFinMsg.Store(conn, &lastMsg)
return q.parseMessage(msg, conn)
}
func checkHeader(msg []byte) bool {
if len(msg) != 8 {
return false
}
for k, v := range msg {
if v != header[k] {
return false
}
}
return true
}
func searchHeader(msg []byte) int {
if len(msg) < 8 {
return 0
}
for k, v := range msg {
find := 0
if v == header[0] {
for k2, v2 := range header {
if msg[k+k2] == v2 {
find++
} else {
break
}
}
if find == 8 {
return k
}
}
}
return -1
}
func bytesMerge(src ...[]byte) []byte {
var buff bytes.Buffer
for _, v := range src {
buff.Write(v)
}
return buff.Bytes()
}
// Restore 获取收到的信息
func (q *StarQueue) Restore() (MsgQueue, error) {
if q.duration.Seconds() == 0 {
q.duration = 86400 * time.Second
}
for {
select {
case <-q.ctx.Done():
return MsgQueue{}, errors.New("Stoped By External Function Call")
case <-time.After(q.duration):
if q.duration != 0 {
return MsgQueue{}, ErrDeadlineExceeded
}
case data, ok := <-q.msgPool:
if !ok {
return MsgQueue{}, os.ErrClosed
}
return data, nil
}
}
}
// RestoreOne 获取收到的一个信息
// 兼容性修改
func (q *StarQueue) RestoreOne() (MsgQueue, error) {
return q.Restore()
}
// Stop 立即停止Restore
func (q *StarQueue) Stop() {
q.cancel()
}
// RestoreDuration Restore最大超时时间
func (q *StarQueue) RestoreDuration(tm time.Duration) {
q.duration = tm
}
func (q *StarQueue) RestoreChan() <-chan MsgQueue {
return q.msgPool
}
+42
View File
@@ -0,0 +1,42 @@
package stario
import (
"fmt"
"testing"
"time"
)
func Test_QueSpeed(t *testing.T) {
que := NewQueueWithCount(0)
stop := make(chan struct{}, 1)
que.RestoreDuration(time.Second * 10)
var count int64
go func() {
for {
select {
case <-stop:
//fmt.Println(count)
return
default:
}
_, err := que.RestoreOne()
if err == nil {
count++
}
}
}()
cp := 0
stoped := time.After(time.Second * 10)
data := que.BuildMessage([]byte("hello"))
for {
select {
case <-stoped:
fmt.Println(count, cp)
stop <- struct{}{}
return
default:
que.ParseMessage(data, "lala")
cp++
}
}
}
+22
View File
@@ -0,0 +1,22 @@
//go:build windows
// +build windows
package stario
import (
"errors"
)
func signal(sigtype rune) error {
//todo: use win32api call signal
switch sigtype {
case 0x03:
return errors.New("SIGNAL SIGINT RECIVED")
case 0x1a:
return errors.New("SIGNAL SIGSTOP RECIVED")
case 0x1c:
return errors.New("SIGNAL SIGQUIT RECIVED")
default:
return nil
}
}
+26
View File
@@ -0,0 +1,26 @@
//go:build !windows
// +build !windows
package stario
import (
"errors"
"os"
"syscall"
)
func signal(sigtype rune) error {
switch sigtype {
case 0x03:
syscall.Kill(os.Getpid(), syscall.SIGINT)
return errors.New("SIGNAL SIGINT RECIVED")
case 0x1a:
syscall.Kill(os.Getpid(), syscall.SIGSTOP)
return errors.New("SIGNAL SIGSTOP RECIVED")
case 0x1c:
syscall.Kill(os.Getpid(), syscall.SIGQUIT)
return errors.New("SIGNAL SIGQUIT RECIVED")
default:
return nil
}
}
+117 -30
View File
@@ -1,55 +1,142 @@
package stario
import (
"fmt"
"sync"
"sync/atomic"
"time"
)
type waitGroupAddMode uint8
const (
waitGroupAddModeStrict waitGroupAddMode = iota
waitGroupAddModeLoose
)
type WaitGroup struct {
wg *sync.WaitGroup
maxCount uint32
allCount uint32
wg sync.WaitGroup
mu sync.Mutex
cond *sync.Cond
initOnce sync.Once
maxCount int
running int
addMode waitGroupAddMode
}
func NewWaitGroup(maxCount int) WaitGroup {
return WaitGroup{wg: &sync.WaitGroup{}, maxCount: uint32(maxCount)}
if maxCount < 0 {
panic("stario: negative max wait count")
}
return WaitGroup{
maxCount: maxCount,
addMode: waitGroupAddModeStrict,
}
}
func (swg *WaitGroup) Add(delta int) {
var Udelta uint32
func (w *WaitGroup) init() {
w.initOnce.Do(func() {
w.cond = sync.NewCond(&w.mu)
})
}
func (w *WaitGroup) Add(delta int) {
w.init()
if delta == 0 {
return
}
if delta < 0 {
Udelta = uint32(-delta - 1)
} else {
Udelta = uint32(delta)
w.release(-delta)
return
}
for {
allC := atomic.LoadUint32(&swg.allCount)
if atomic.LoadUint32(&swg.maxCount) == 0 || atomic.LoadUint32(&swg.maxCount) >= allC+uint32(delta) {
if delta < 0 {
atomic.AddUint32(&swg.allCount, ^uint32(Udelta))
} else {
atomic.AddUint32(&swg.allCount, uint32(Udelta))
}
break
w.acquire(delta)
}
func (w *WaitGroup) acquire(delta int) {
w.mu.Lock()
defer w.mu.Unlock()
if w.maxCount <= 0 {
w.running += delta
w.wg.Add(delta)
return
}
if delta == 1 {
w.wg.Add(1)
for w.maxCount > 0 && w.running >= w.maxCount {
w.cond.Wait()
}
time.Sleep(time.Microsecond)
w.running++
return
}
swg.wg.Add(delta)
if w.running+delta > w.maxCount {
if w.addMode == waitGroupAddModeStrict {
panic(fmt.Sprintf("stario: WaitGroup.Add(%d) exceeds max limit %d with %d running", delta, w.maxCount, w.running))
}
w.maxCount = w.running + delta
}
w.running += delta
w.wg.Add(delta)
}
func (swg *WaitGroup) Done() {
swg.Add(-1)
func (w *WaitGroup) release(delta int) {
w.mu.Lock()
defer w.mu.Unlock()
if delta > w.running {
panic(fmt.Sprintf("stario: WaitGroup.Done releases %d tasks but only %d running", delta, w.running))
}
w.wg.Add(-delta)
w.running -= delta
w.cond.Broadcast()
}
func (swg *WaitGroup) Wait() {
swg.wg.Wait()
func (w *WaitGroup) Done() {
w.Add(-1)
}
func (swg *WaitGroup) GetMaxWaitNum() int {
return int(atomic.LoadUint32(&swg.maxCount))
func (w *WaitGroup) Go(fn func()) {
w.Add(1)
go func() {
defer w.Done()
fn()
}()
}
func (swg *WaitGroup) SetMaxWaitNum(num int) {
atomic.AddUint32(&swg.maxCount, uint32(num))
func (w *WaitGroup) Wait() {
w.init()
w.wg.Wait()
}
func (w *WaitGroup) GetMaxWaitNum() int {
w.init()
w.mu.Lock()
defer w.mu.Unlock()
return w.maxCount
}
func (w *WaitGroup) SetMaxWaitNum(num int) {
if num < 0 {
panic("stario: negative max wait count")
}
w.init()
w.mu.Lock()
w.maxCount = num
w.mu.Unlock()
w.cond.Broadcast()
}
func (w *WaitGroup) SetStrictAddMode(strict bool) {
w.init()
w.mu.Lock()
if strict {
w.addMode = waitGroupAddModeStrict
} else {
w.addMode = waitGroupAddModeLoose
}
w.mu.Unlock()
w.cond.Broadcast()
}
func (w *WaitGroup) StrictAddMode() bool {
w.init()
w.mu.Lock()
defer w.mu.Unlock()
return w.addMode == waitGroupAddModeStrict
}
+171
View File
@@ -0,0 +1,171 @@
package stario
import (
"testing"
"time"
)
func TestWaitGroupAddBlocksAtLimit(t *testing.T) {
wg := NewWaitGroup(1)
wg.Add(1)
unblocked := make(chan struct{})
go func() {
wg.Add(1)
close(unblocked)
wg.Done()
}()
select {
case <-unblocked:
t.Fatal("Add(1) should block while running count is at the limit")
case <-time.After(30 * time.Millisecond):
}
wg.Done()
select {
case <-unblocked:
case <-time.After(200 * time.Millisecond):
t.Fatal("Add(1) did not resume after capacity was released")
}
wg.Wait()
}
func TestWaitGroupStrictAddPanicsWhenBatchExceedsLimit(t *testing.T) {
wg := NewWaitGroup(1)
wg.Add(1)
defer wg.Done()
defer func() {
if recover() == nil {
t.Fatal("expected Add(2) to panic in strict mode when it exceeds the limit")
}
}()
wg.Add(2)
}
func TestWaitGroupLooseAddExpandsLimit(t *testing.T) {
wg := NewWaitGroup(1)
wg.SetStrictAddMode(false)
wg.Add(1)
wg.Add(2)
if got := wg.GetMaxWaitNum(); got != 3 {
t.Fatalf("unexpected max count: got %d want 3", got)
}
for i := 0; i < 3; i++ {
wg.Done()
}
wg.Wait()
}
func TestWaitGroupSetMaxWaitNumOverridesValue(t *testing.T) {
wg := NewWaitGroup(1)
wg.SetMaxWaitNum(4)
if got := wg.GetMaxWaitNum(); got != 4 {
t.Fatalf("unexpected max count: got %d want 4", got)
}
}
func TestWaitGroupSetMaxWaitNumZeroUnblocksWaitingAdd(t *testing.T) {
wg := NewWaitGroup(1)
wg.Add(1)
unblocked := make(chan struct{})
go func() {
wg.Add(1)
close(unblocked)
wg.Done()
}()
select {
case <-unblocked:
t.Fatal("Add(1) should still be blocked before removing the limit")
case <-time.After(30 * time.Millisecond):
}
wg.SetMaxWaitNum(0)
select {
case <-unblocked:
case <-time.After(200 * time.Millisecond):
t.Fatal("Add(1) did not resume after the limit was removed")
}
wg.Done()
wg.Wait()
}
func TestWaitGroupGoWaitsForCompletion(t *testing.T) {
wg := NewWaitGroup(1)
done := make(chan struct{})
wg.Go(func() {
close(done)
})
select {
case <-done:
case <-time.After(200 * time.Millisecond):
t.Fatal("Go(fn) did not run the function")
}
finished := make(chan struct{})
go func() {
wg.Wait()
close(finished)
}()
select {
case <-finished:
case <-time.After(200 * time.Millisecond):
t.Fatal("Wait did not observe Go(fn) completion")
}
}
func TestWaitGroupGoRespectsLimit(t *testing.T) {
wg := NewWaitGroup(1)
releaseFirst := make(chan struct{})
firstStarted := make(chan struct{})
secondStarted := make(chan struct{})
secondReturned := make(chan struct{})
wg.Go(func() {
close(firstStarted)
<-releaseFirst
})
select {
case <-firstStarted:
case <-time.After(200 * time.Millisecond):
t.Fatal("first Go(fn) did not start")
}
go func() {
wg.Go(func() {
close(secondStarted)
})
close(secondReturned)
}()
select {
case <-secondReturned:
t.Fatal("second Go(fn) should block while the limit is full")
case <-time.After(30 * time.Millisecond):
}
close(releaseFirst)
select {
case <-secondStarted:
case <-time.After(200 * time.Millisecond):
t.Fatal("second Go(fn) did not start after capacity was released")
}
wg.Wait()
}