func main() {
kingpin.Parse()
msgChan, err := buildStream()
if err != nil || msgChan == nil {
kingpin.FatalUsage("Failed to create input stream: " + err.Error())
}
// Disable regular logging, cuz bonjour logs an error that isn't, and it's
// confusing.
dontlog.SetOutput(NullWriter{})
if *host == "" {
*host = getLocalIP()
}
if *port == 0 {
// TODO: run a server that sends the same text on TCP conn
*port = 45897
}
serv, err := service.New(*host, *port, msgChan)
if err != nil || serv == nil {
log.Panic("Failed to create service", "err", err)
}
// Watch for signal to clean up before we exit
signals := make(chan os.Signal, 1)
signal.Notify(signals, os.Interrupt, os.Kill)
<-signals
log.Info("Shutting down")
serv.Stop()
log.Info("Stopped service")
}
func (s *Service) Stop() {
defer s.wg.Wait()
log.Info("Stopping service")
close(s.stop)
// TCPListener might still be blocked, so open a conn to it ourselves
// to free it up.
if conn, err := net.DialTCP("tcp", nil, s.addr); err != nil {
log.Error("Failed to close down TCP listener", "err", err)
} else {
conn.Close()
log.Info("Shut down TCP listener")
}
}
func New(host string, port int, msgChan <-chan string) (*Service, error) {
s := &Service{
host: host,
port: port,
messages: msgChan,
stop: make(chan struct{}),
}
if addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", host, port)); err != nil {
return nil, err
} else {
s.addr = addr
}
listener, err := net.ListenTCP("tcp", s.addr)
if err != nil {
return nil, err
}
log.Info("Listening for TCP connections", "address", s.addr)
s.wg.Add(1)
go s.serveTCP(listener)
s.wg.Add(1)
go s.start()
return s, nil
}
func CSVField(fieldIndex int, source DataSource) MessageBuilder {
c := make(chan []string)
go func() {
for {
reader := csv.NewReader(strings.NewReader(<-source))
reader.TrimLeadingSpace = true
reader.FieldsPerRecord = -1
var values []string
// Ignore first line - comment
reader.Read()
for {
if record, err := reader.Read(); err == io.EOF {
break
} else if err != nil {
log.Error("Error reading CSV data", "err", err)
} else if len(record) <= fieldIndex {
log.Error("Error reading CSV data: index out of bounds", "# fields in record", record, "index", fieldIndex)
} else {
values = append(values, record[fieldIndex])
}
}
if len(values) == 0 {
log.Error("Didn't get any values from CSV")
} else {
log.Info("Read CSV data", "num values", len(values))
c <- values
}
}
}()
return c
}
func (s *Service) stopBonjour(bonj *bonjour.Server) {
if bonj == nil {
return
}
s.wg.Add(1)
go func() {
defer s.wg.Done()
log.Info("Shutting down bonjour service")
bonj.Shutdown()
// I guess bonjour wants us to wait some unspecied
// amount? This is what blocking or channels are for :/
waitTime := time.Second * 5
log.Info("Waiting for bonjour service to clean itself up", "waitTime", waitTime)
time.Sleep(waitTime)
}()
}
func (s *Service) start() {
defer s.wg.Done()
var bonj *bonjour.Server
defer func(b **bonjour.Server) {
s.stopBonjour(*b)
}(&bonj)
for {
select {
case msg := <-s.messages:
if msg != s.currentMsg {
s.currentMsg = msg
s.stopBonjour(bonj)
bonj = nil
if msg != "" {
log.Info("Registering service", "name", msg, "host", s.host, "port", s.port)
var err error
bonj, err = bonjour.RegisterProxy(
msg,
"_afpovertcp._tcp", "local",
s.port, s.host, s.host,
nil, nil)
if err != nil || bonj == nil {
log.Error("Failed to register service with bonjour", "err", err)
bonj = nil
}
}
}
case <-s.stop:
log.Info("Ending bonjour-updating routine")
return
}
}
}
func buildStream() (<-chan string, error) {
var source inputs.DataSource
var err error
if *url != "" {
log.Info("Reading messages from a url", "url", *url)
source, err = inputs.Download(*url)
if err != nil {
return nil, err
}
} else if *file != "" {
log.Info("Reading messages from file", "file", *file)
source, err = inputs.FileWatcher(*file)
if err != nil {
return nil, err
}
} else if len(*say) > 0 {
log.Info("Using a static message", "msg", *say)
source, err = inputs.StaticText(strings.Join(*say, " "))
if err != nil {
return nil, err
}
}
if source == nil {
return nil, errors.New("No source of data specified")
}
var builder inputs.MessageBuilder
if *csvField >= 0 {
log.Info("Iterating csv values", "field", *csvField)
builder = inputs.CSVField(*csvField, source)
} else if *words {
log.Info("Iterating words")
builder = inputs.WordGroups(source)
} else {
log.Info("Iterating lines")
builder = inputs.Lines(source)
}
var chooser inputs.MessageChooser
if *random {
log.Info("Randomizing order")
chooser = inputs.RandomMessageChooser(builder)
} else {
log.Info("Sequential order")
chooser = inputs.SequentialMessageChooser(builder)
}
msgChan := inputs.RateLimit(*interval,
inputs.LimitSize(40,
inputs.Cleanup(
chooser)))
// Optional filters
if *prefix != "" {
msgChan = inputs.Prefix(*prefix, msgChan)
}
if *leet {
msgChan = inputs.LeetSpeak(msgChan)
} else if *mixedCase {
msgChan = inputs.MixedCase(msgChan)
} else if *lower {
msgChan = inputs.LowerCase(msgChan)
} else if *upper {
msgChan = inputs.UpperCase(msgChan)
}
return msgChan, nil
}