Rob Glew
7 years ago
parent
06defd68df
commit
b5a7a93853
2 changed files with 0 additions and 280 deletions
@ -1,117 +0,0 @@ |
||||
package main |
||||
|
||||
import ( |
||||
"fmt" |
||||
"sync" |
||||
) |
||||
|
||||
// An interface which represents a job to be done by the pool
|
||||
type Job interface { |
||||
Run() // Start the job
|
||||
Abort() // Abort any work that needs to be completed and close the DoneChannel
|
||||
DoneChannel() chan struct{} // Returns a channel that is closed when the job is done
|
||||
} |
||||
|
||||
// An interface which represents a pool of workers doing jobs
|
||||
type JobPool struct { |
||||
MaxThreads int |
||||
|
||||
jobQueue chan Job |
||||
jobQueueDone chan struct{} |
||||
jobQueueAborted chan struct{} |
||||
jobQueueShutDown chan struct{} |
||||
} |
||||
|
||||
func NewJobPool(maxThreads int) *JobPool { |
||||
q := JobPool{ |
||||
MaxThreads: maxThreads, |
||||
jobQueue: make(chan Job), |
||||
jobQueueDone: make(chan struct{}), // Closing will shut down workers and reject any incoming work
|
||||
jobQueueAborted: make(chan struct{}), // Closing tells workers to abort
|
||||
jobQueueShutDown: make(chan struct{}), // Closed when all workers are shut down
|
||||
} |
||||
|
||||
return &q |
||||
} |
||||
|
||||
func (q *JobPool) RunJob(j Job) error { |
||||
select { |
||||
case <-q.jobQueueDone: |
||||
return fmt.Errorf("job queue closed") |
||||
default: |
||||
} |
||||
|
||||
q.jobQueue <- j |
||||
return nil |
||||
} |
||||
|
||||
func (q *JobPool) Run() { |
||||
if q.MaxThreads > 0 { |
||||
// Create pool of routines that read from the queue and run jobs
|
||||
var w sync.WaitGroup |
||||
for i := 0; i < q.MaxThreads; i++ { |
||||
w.Add(1) |
||||
go func() { |
||||
defer w.Done() |
||||
for { |
||||
select { |
||||
case job := <-q.jobQueue: |
||||
go func() { |
||||
job.Run() |
||||
}() |
||||
|
||||
select { |
||||
case <-job.DoneChannel(): // The job finishes normally
|
||||
case <-q.jobQueueAborted: // We have to abort the job
|
||||
job.Abort() // Tell the job to abort
|
||||
<-job.DoneChannel() // Wait for the job to abort
|
||||
} |
||||
case <-q.jobQueueDone: |
||||
// We're done and out of jobs, quit
|
||||
close(q.jobQueueShutDown) // Flag that the workers quit
|
||||
return |
||||
} |
||||
} |
||||
}() |
||||
} |
||||
|
||||
w.Wait() // Wait for workers to quit
|
||||
close(q.jobQueueShutDown) // Flag that all workers quit
|
||||
} else { |
||||
// Create a thread any time we pull something out of the job queue
|
||||
for { |
||||
select { |
||||
case job := <-q.jobQueue: |
||||
go func() { |
||||
go func() { |
||||
job.Run() |
||||
}() |
||||
|
||||
select { |
||||
case <-job.DoneChannel(): // The job finishes normally
|
||||
case <-q.jobQueueAborted: // We have to abort the job
|
||||
job.Abort() // Tell the job to abort
|
||||
<-job.DoneChannel() // Wait for the job to abort
|
||||
} |
||||
}() |
||||
case <-q.jobQueueDone: |
||||
close(q.jobQueueShutDown) // Flag that the workers quit
|
||||
return |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
func (q *JobPool) Abort() { |
||||
close(q.jobQueueDone) // Stop accepting jobs and tell the workers to quit
|
||||
close(q.jobQueueAborted) // Tell the workers to abort
|
||||
<-q.jobQueueShutDown // Wait for all the workers to shut down
|
||||
close(q.jobQueue) // Clean up the job queue
|
||||
} |
||||
|
||||
func (q *JobPool) CompleteAndClose() { |
||||
close(q.jobQueueDone) // Stop accepting jobs and tell the workers to quit
|
||||
<-q.jobQueueShutDown // Wait for all the workers to shut down
|
||||
close(q.jobQueue) // Clean up the job queue
|
||||
close(q.jobQueueAborted) // Clean up abort channel
|
||||
} |
||||
@ -1,163 +0,0 @@ |
||||
package main |
||||
|
||||
import ( |
||||
"errors" |
||||
"flag" |
||||
"fmt" |
||||
"io/ioutil" |
||||
"log" |
||||
"net" |
||||
"os" |
||||
"os/signal" |
||||
"strings" |
||||
"syscall" |
||||
"time" |
||||
) |
||||
|
||||
var logBanner string = ` |
||||
======================================== |
||||
PUPPYSTARTEDPUPPYSTARTEDPUPPYSTARTEDPUPP |
||||
.--. .---. |
||||
/:. '. .' .. '._.---. |
||||
/:::-. \.-"""-;' .-:::. .::\
|
||||
/::'| '\/ _ _ \' '\:' ::::| |
||||
__.' | / (o|o) \ ''. ':/ |
||||
/ .:. / | ___ | '---' |
||||
| ::::' /: (._.) .:\
|
||||
\ .=' |:' :::| |
||||
'""' \ .-. ':/ |
||||
'---'|I|'---' |
||||
'-' |
||||
PUPPYSTARTEDPUPPYSTARTEDPUPPYSTARTEDPUPP |
||||
======================================== |
||||
` |
||||
|
||||
type listenArg struct { |
||||
Type string |
||||
Addr string |
||||
} |
||||
|
||||
func quitErr(msg string) { |
||||
os.Stderr.WriteString(msg) |
||||
os.Stderr.WriteString("\n") |
||||
os.Exit(1) |
||||
} |
||||
|
||||
func checkErr(err error) { |
||||
if err != nil { |
||||
quitErr(err.Error()) |
||||
} |
||||
} |
||||
|
||||
func parseListenString(lstr string) (*listenArg, error) { |
||||
args := strings.SplitN(lstr, ":", 2) |
||||
if len(args) != 2 { |
||||
return nil, errors.New("invalid listener. Must be in the form of \"tye:addr\"") |
||||
} |
||||
argStruct := &listenArg{ |
||||
Type: strings.ToLower(args[0]), |
||||
Addr: args[1], |
||||
} |
||||
if argStruct.Type != "tcp" && argStruct.Type != "unix" { |
||||
return nil, fmt.Errorf("invalid listener type: %s", argStruct.Type) |
||||
} |
||||
return argStruct, nil |
||||
} |
||||
|
||||
func unixAddr() string { |
||||
return fmt.Sprintf("%s/proxy.%d.%d.sock", os.TempDir(), os.Getpid(), time.Now().UnixNano()) |
||||
} |
||||
|
||||
var mln net.Listener |
||||
var logger *log.Logger |
||||
|
||||
func cleanup() { |
||||
if mln != nil { |
||||
mln.Close() |
||||
} |
||||
} |
||||
|
||||
var MainLogger *log.Logger |
||||
|
||||
func main() { |
||||
defer cleanup() |
||||
// Handle signals
|
||||
sigc := make(chan os.Signal, 1) |
||||
signal.Notify(sigc, os.Interrupt, os.Kill, syscall.SIGTERM) |
||||
go func() { |
||||
<-sigc |
||||
if logger != nil { |
||||
logger.Println("Caught signal. Cleaning up.") |
||||
} |
||||
cleanup() |
||||
os.Exit(0) |
||||
}() |
||||
|
||||
msgListenStr := flag.String("msglisten", "", "Listener for the message handler. Examples: \"tcp::8080\", \"tcp:127.0.0.1:8080\", \"unix:/tmp/foobar\"") |
||||
// storageFname := flag.String("storage", "", "Datafile to use for storage")
|
||||
// inMemStorage := flag.Bool("inmem", false, "Set flag to store messages in memroy rather than use a datafile")
|
||||
autoListen := flag.Bool("msgauto", false, "Automatically pick and open a unix or tcp socket for the message listener") |
||||
debugFlag := flag.Bool("dbg", false, "Enable debug logging") |
||||
flag.Parse() |
||||
|
||||
if *debugFlag { |
||||
logfile, err := os.OpenFile("log.log", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) |
||||
checkErr(err) |
||||
logger = log.New(logfile, "[*] ", log.Lshortfile) |
||||
} else { |
||||
logger = log.New(ioutil.Discard, "[*] ", log.Lshortfile) |
||||
log.SetFlags(0) |
||||
} |
||||
MainLogger = logger |
||||
|
||||
// Parse arguments to structs
|
||||
if *msgListenStr == "" && *autoListen == false { |
||||
quitErr("message listener address or `--msgauto` required") |
||||
} |
||||
if *msgListenStr != "" && *autoListen == true { |
||||
quitErr("only one of listener address or `--msgauto` can be used") |
||||
} |
||||
|
||||
// Create the message listener
|
||||
var listenStr string |
||||
if *msgListenStr != "" { |
||||
msgAddr, err := parseListenString(*msgListenStr) |
||||
checkErr(err) |
||||
if msgAddr.Type == "tcp" { |
||||
var err error |
||||
mln, err = net.Listen("tcp", msgAddr.Addr) |
||||
checkErr(err) |
||||
} else if msgAddr.Type == "unix" { |
||||
var err error |
||||
mln, err = net.Listen("unix", msgAddr.Addr) |
||||
checkErr(err) |
||||
} else { |
||||
quitErr("unsupported listener type:" + msgAddr.Type) |
||||
} |
||||
listenStr = fmt.Sprintf("%s:%s", msgAddr.Type, msgAddr.Addr) |
||||
} else { |
||||
fpath := unixAddr() |
||||
ulisten, err := net.Listen("unix", fpath) |
||||
if err == nil { |
||||
mln = ulisten |
||||
listenStr = fmt.Sprintf("unix:%s", fpath) |
||||
} else { |
||||
tcplisten, err := net.Listen("tcp", "127.0.0.1:0") |
||||
if err != nil { |
||||
quitErr("unable to open any messaging ports") |
||||
} |
||||
mln = tcplisten |
||||
listenStr = fmt.Sprintf("tcp:%s", tcplisten.Addr().String()) |
||||
} |
||||
} |
||||
|
||||
// Set up the intercepting proxy
|
||||
iproxy := NewInterceptingProxy(logger) |
||||
iproxy.AddHTTPHandler("puppy", WebUIHandler) |
||||
|
||||
// Create a message server and have it serve for the iproxy
|
||||
mserv := NewProxyMessageListener(logger, iproxy) |
||||
logger.Print(logBanner) |
||||
fmt.Println(listenStr) |
||||
mserv.Serve(mln) // serve until killed
|
||||
} |
||||
Loading…
Reference in new issue