func RecMsgFromPeers(s *zmq.Socket, ser Servernode) {
var e Envelope
for {
msg, err := s.RecvMessage(0)
a1 := strings.Split(msg[0], "&{")
a2 := strings.Split(a1[1], " ")
pid, _ := strconv.Atoi(a2[0])
mid, _ := strconv.Atoi(a2[1])
msgid := int64(mid)
text := ""
for b := 2; b < len(a2); b++ {
text = text + " " + a2[b]
}
text1 := text[1:]
msgtext := strings.Split(text1, "}")[0]
e.Pid = pid
e.MsgId = msgid
e.Msg = msgtext
go QueueInbox(ser, e, ser.Inbox())
if err != nil {
break
}
s.SendMessage(msg)
}
}
func (c *WorkerConnection) run(socket *zmq.Socket) {
defer socket.Close()
socket.SetRcvtimeo(c.activeTimeout)
for {
select {
case <-c.quit:
close(c.wait)
return
case request := <-c.requestChannel:
if _, err := socket.SendMessage(request); err != nil {
log.Println("Failed to send request:", err)
}
case request := <-c.connectionChan:
go c.handleConnectionRequest(socket, request)
default:
message, err := socket.RecvMessageBytes(0)
if err != nil {
// Needed to yield to goroutines when GOMAXPROCS is 1.
// Note: The 1.3 preemptive scheduler doesn't seem to work here,
// so this is still required.
runtime.Gosched()
break
}
socket.SetRcvtimeo(c.activeTimeout)
go c.handleResponse(message)
}
}
}
func reflector(t *testing.T, s *zmq.Socket, n int, wg *sync.WaitGroup) chan bool {
s.SetRcvtimeo(time.Second)
s.SetSndtimeo(time.Second)
svcrouter.Barrier()
wg.Add(1)
c := make(chan bool)
go func() {
/* this is a nasty hack but is OK for a test program I guess */
for {
if msg, err := s.RecvMessageBytes(0); err == nil {
svcrouter.DumpMsg(fmt.Sprintf("reflector %d", n), msg)
if _, err := s.SendMessage(msg); err != nil {
wg.Done()
t.Fatal(err)
return
}
}
select {
case <-c:
wg.Done()
return
default:
}
}
}()
return c
}
// Echo service
func echo(socket *zmq.Socket) (err error) {
msg, err := socket.RecvMessage(0)
if err != nil {
return
}
_, err = socket.SendMessage(msg)
return
}
func openDoor(sp gpio.DirectPinDriver, publisher *zmq.Socket) {
sp.DigitalWrite(1)
publisher.SendMessage("door.state.unlock", "Door Unlocked")
gobot.After(5*time.Second, func() {
sp.DigitalWrite(0)
publisher.SendMessage("door.state.lock", "Door Locked")
})
}
func (k *kernelRunner) sendMessage(msg *message, ids []string, socket *zmq.Socket) error {
parts, err := serializeMessage(msg, ids, []byte(k.connInfo.Key))
if err != nil {
return fmt.Errorf("serializing message: %v", err)
}
if _, err := socket.SendMessage(parts); err != nil {
return fmt.Errorf("sending message: %v", err)
}
return nil
}
/* {{{ func RequestAndReply(soc *zmq.Socket, msg interface{}) (reply []string, err error) {
* 支持超时时间,但要注意如果一个socket失败之后，因为req严格同步,连接池模式下这个socket最好销毁
*/
func RequestAndReply(soc *zmq.Socket, timeout time.Duration, msg ...interface{}) (reply []string, err error) {
poller := zmq.NewPoller()
poller.Add(soc, zmq.POLLIN)
if _, err := soc.SendMessage(msg...); err != nil {
return nil, err
}
if sockets, err := poller.Poll(timeout); err != nil {
return nil, err
} else if len(sockets) == 1 {
return soc.RecvMessage(zmq.DONTWAIT)
} else {
return nil, fmt.Errorf("time out!")
}
return
}
func NewReplyHandler() chan<- *ReplyMessage {
replyChan := make(chan *ReplyMessage, replyChanCapacity)
pushSockets := make(map[string]*zmq4.Socket)
go func() {
for replyMessage := range replyChan {
marshalledReply, err := json.Marshal(replyMessage.Content)
if err != nil {
fog.Error("unable to marshall reply %s %s", replyMessage, err)
continue
}
var pushSocket *zmq4.Socket
var ok bool
pushSocket, ok = pushSockets[replyMessage.ClientAddress]
if !ok {
fog.Info("creating PUSH socket to %s", replyMessage.ClientAddress)
if pushSocket, err = createPushSocket(); err != nil {
fog.Error("Unable to create PUSH socket for %s %s",
replyMessage.ClientAddress, err)
continue
}
if err = pushSocket.Connect(replyMessage.ClientAddress); err != nil {
fog.Error("Unable to Connect PUSH socket to %s %s",
replyMessage.ClientAddress, err)
pushSocket.Close()
continue
}
pushSockets[replyMessage.ClientAddress] = pushSocket
}
if _, err = pushSocket.SendMessage(marshalledReply); err != nil {
fog.Error("pushSocket SendMessage to %s failed %s",
replyMessage.ClientAddress, err)
pushSocket.Close()
delete(pushSockets, replyMessage.ClientAddress)
continue
}
}
}()
return replyChan
}
// The send method sends a multi-frame key-value message to a socket.
func (kvmsg *Kvmsg) Send(socket *zmq.Socket) (err error) {
//fmt.Printf("Send to %s: %q\n", socket, kvmsg.frame)
_, err = socket.SendMessage(kvmsg.frame)
return
}
func TestMultipleContexts(t *testing.T) {
chQuit := make(chan interface{})
chErr := make(chan error, 2)
needQuit := false
var sock1, sock2, serv1, serv2 *zmq.Socket
var serv_ctx1, serv_ctx2, ctx1, ctx2 *zmq.Context
var err error
defer func() {
if needQuit {
chQuit <- true
chQuit <- true
<-chErr
<-chErr
}
for _, s := range []*zmq.Socket{sock1, sock2, serv1, serv2} {
if s != nil {
s.SetLinger(0)
s.Close()
}
}
for _, c := range []*zmq.Context{serv_ctx1, serv_ctx2, ctx1, ctx2} {
if c != nil {
c.Term()
}
}
}()
addr1 := "tcp://127.0.0.1:9997"
addr2 := "tcp://127.0.0.1:9998"
serv_ctx1, err = zmq.NewContext()
if err != nil {
t.Fatal("NewContext:", err)
}
serv1, err = serv_ctx1.NewSocket(zmq.REP)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = serv1.Bind(addr1)
if err != nil {
t.Fatal("Bind:", err)
}
serv_ctx2, err = zmq.NewContext()
if err != nil {
t.Fatal("NewContext:", err)
}
serv2, err = serv_ctx2.NewSocket(zmq.REP)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = serv2.Bind(addr2)
if err != nil {
t.Fatal("Bind:", err)
}
new_service := func(sock *zmq.Socket, addr string) {
socket_handler := func(state zmq.State) error {
msg, err := sock.RecvMessage(0)
if err != nil {
return err
}
_, err = sock.SendMessage(addr, msg)
return err
}
quit_handler := func(interface{}) error {
return errors.New("quit")
}
reactor := zmq.NewReactor()
reactor.AddSocket(sock, zmq.POLLIN, socket_handler)
reactor.AddChannel(chQuit, 1, quit_handler)
err = reactor.Run(100 * time.Millisecond)
chErr <- err
}
go new_service(serv1, addr1)
go new_service(serv2, addr2)
needQuit = true
time.Sleep(time.Second)
// default context
sock1, err = zmq.NewSocket(zmq.REQ)
if err != nil {
t.Fatal("NewSocket:", err)
}
sock2, err = zmq.NewSocket(zmq.REQ)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = sock1.Connect(addr1)
if err != nil {
t.Fatal("sock1.Connect:", err)
}
err = sock2.Connect(addr2)
if err != nil {
t.Fatal("sock2.Connect:", err)
}
_, err = sock1.SendMessage(addr1)
if err != nil {
t.Fatal("sock1.SendMessage:", err)
}
_, err = sock2.SendMessage(addr2)
if err != nil {
t.Fatal("sock2.SendMessage:", err)
}
msg, err := sock1.RecvMessage(0)
expected := []string{addr1, addr1}
if err != nil || !arrayEqual(msg, expected) {
t.Errorf("sock1.RecvMessage: expected %v %v, got %v %v", nil, expected, err, msg)
}
msg, err = sock2.RecvMessage(0)
expected = []string{addr2, addr2}
if err != nil || !arrayEqual(msg, expected) {
t.Errorf("sock2.RecvMessage: expected %v %v, got %v %v", nil, expected, err, msg)
}
err = sock1.Close()
sock1 = nil
if err != nil {
t.Fatal("sock1.Close:", err)
}
err = sock2.Close()
sock2 = nil
if err != nil {
t.Fatal("sock2.Close:", err)
}
// non-default contexts
ctx1, err = zmq.NewContext()
if err != nil {
t.Fatal("NewContext:", err)
}
ctx2, err = zmq.NewContext()
if err != nil {
t.Fatal("NewContext:", err)
}
sock1, err = ctx1.NewSocket(zmq.REQ)
if err != nil {
t.Fatal("ctx1.NewSocket:", err)
}
sock2, err = ctx2.NewSocket(zmq.REQ)
if err != nil {
t.Fatal("ctx2.NewSocket:", err)
}
err = sock1.Connect(addr1)
if err != nil {
t.Fatal("sock1.Connect:", err)
}
err = sock2.Connect(addr2)
if err != nil {
t.Fatal("sock2.Connect:", err)
}
_, err = sock1.SendMessage(addr1)
if err != nil {
t.Fatal("sock1.SendMessage:", err)
}
_, err = sock2.SendMessage(addr2)
if err != nil {
t.Fatal("sock2.SendMessage:", err)
}
msg, err = sock1.RecvMessage(0)
expected = []string{addr1, addr1}
if err != nil || !arrayEqual(msg, expected) {
t.Errorf("sock1.RecvMessage: expected %v %v, got %v %v", nil, expected, err, msg)
}
msg, err = sock2.RecvMessage(0)
expected = []string{addr2, addr2}
if err != nil || !arrayEqual(msg, expected) {
t.Errorf("sock2.RecvMessage: expected %v %v, got %v %v", nil, expected, err, msg)
}
err = sock1.Close()
sock1 = nil
if err != nil {
t.Fatal("sock1.Close:", err)
}
err = sock2.Close()
sock2 = nil
if err != nil {
t.Fatal("sock2.Close:", err)
}
err = ctx1.Term()
ctx1 = nil
if err != nil {
t.Fatal("ctx1.Term", nil)
}
err = ctx2.Term()
ctx1 = nil
if err != nil {
t.Fatal("ctx2.Term", nil)
}
needQuit = false
for i := 0; i < 2; i++ {
// close(chQuit) doesn't work because the reactor removes closed channels, instead of acting on them
chQuit <- true
err = <-chErr
if err.Error() != "quit" {
t.Errorf("Expected error value quit, got %v", err)
}
}
}
func TestSecurityPlain(t *testing.T) {
time.Sleep(100 * time.Millisecond)
var handler, server, client *zmq.Socket
defer func() {
for _, s := range []*zmq.Socket{handler} {
if s != nil {
s.SetLinger(0)
s.Close()
}
}
}()
handler, err := zmq.NewSocket(zmq.REP)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = handler.Bind("inproc://zeromq.zap.01")
if err != nil {
t.Fatal("handler.Bind:", err)
}
doHandler := func(state zmq.State) error {
msg, err := handler.RecvMessage(0)
if err != nil {
return err // Terminating
}
version := msg[0]
sequence := msg[1]
// domain := msg[2]
// address := msg[3]
identity := msg[4]
mechanism := msg[5]
username := msg[6]
password := msg[7]
if version != "1.0" {
return errors.New("version != 1.0")
}
if mechanism != "PLAIN" {
return errors.New("mechanism != PLAIN")
}
if identity != "IDENT" {
return errors.New("identity != IDENT")
}
if username == "admin" && password == "password" {
handler.SendMessage(version, sequence, "200", "OK", "anonymous", "")
} else {
handler.SendMessage(version, sequence, "400", "Invalid username or password", "", "")
}
return nil
}
doQuit := func(i interface{}) error {
err := handler.Close()
if err != nil {
t.Error("handler.Close:", err)
}
return errors.New("Quit")
}
quit := make(chan interface{})
reactor := zmq.NewReactor()
reactor.AddSocket(handler, zmq.POLLIN, doHandler)
reactor.AddChannel(quit, 0, doQuit)
go func() {
reactor.Run(100 * time.Millisecond)
quit <- true
}()
defer func() {
quit <- true
<-quit
close(quit)
}()
// Server socket will accept connections
server, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket", err)
}
err = server.SetIdentity("IDENT")
if err != nil {
t.Fatal("server.SetIdentity:", err)
}
err = server.SetPlainServer(1)
if err != nil {
t.Fatal("server.SetPlainServer(1):", err)
}
err = server.Bind("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("server.Bind")
}
// Check PLAIN security with correct username/password
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetPlainUsername("admin")
if err != nil {
t.Fatal("client.SetPlainUsername:"******"password")
if err != nil {
t.Fatal("client.SetPlainPassword:"******"tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
msg, err := bounce(server, client)
if err != nil {
t.Error(msg, err)
}
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
// Check PLAIN security with badly configured client (as_server)
// This will be caught by the plain_server class, not passed to ZAP
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
client.SetPlainServer(1)
if err != nil {
t.Fatal("client.SetPlainServer(1):", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
err = server.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("server.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
client.SetLinger(0)
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
err = server.Close()
server = nil
if err != nil {
t.Fatal("server.Close:", err)
}
}
func TestSecurityNull(t *testing.T) {
time.Sleep(100 * time.Millisecond)
var handler, server, client *zmq.Socket
defer func() {
for _, s := range []*zmq.Socket{handler} {
if s != nil {
s.SetLinger(0)
s.Close()
}
}
}()
handler, err := zmq.NewSocket(zmq.REP)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = handler.Bind("inproc://zeromq.zap.01")
if err != nil {
t.Fatal("handler.Bind:", err)
}
doHandler := func(state zmq.State) error {
msg, err := handler.RecvMessage(0)
if err != nil {
return err // Terminating
}
version := msg[0]
sequence := msg[1]
domain := msg[2]
// address := msg[3]
// identity := msg[4]
mechanism := msg[5]
if version != "1.0" {
return errors.New("version != 1.0")
}
if mechanism != "NULL" {
return errors.New("mechanism != NULL")
}
if domain == "TEST" {
handler.SendMessage(version, sequence, "200", "OK", "anonymous", "")
} else {
handler.SendMessage(version, sequence, "400", "BAD DOMAIN", "", "")
}
return nil
}
doQuit := func(i interface{}) error {
err := handler.Close()
handler = nil
if err != nil {
t.Error("handler.Close:", err)
}
return errors.New("Quit")
}
quit := make(chan interface{})
reactor := zmq.NewReactor()
reactor.AddSocket(handler, zmq.POLLIN, doHandler)
reactor.AddChannel(quit, 0, doQuit)
go func() {
reactor.Run(100 * time.Millisecond)
quit <- true
}()
defer func() {
quit <- true
<-quit
close(quit)
}()
// We bounce between a binding server and a connecting client
server, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
// We first test client/server with no ZAP domain
// Libzmq does not call our ZAP handler, the connect must succeed
err = server.Bind("tcp://127.0.0.1:9683")
if err != nil {
t.Fatal("server.Bind:", err)
}
err = client.Connect("tcp://127.0.0.1:9683")
if err != nil {
t.Fatal("client.Connect:", err)
}
msg, err := bounce(server, client)
if err != nil {
t.Error(msg, err)
}
server.Unbind("tcp://127.0.0.1:9683")
client.Disconnect("tcp://127.0.0.1:9683")
// Now define a ZAP domain for the server; this enables
// authentication. We're using the wrong domain so this test
// must fail.
err = server.SetZapDomain("WRONG")
if err != nil {
t.Fatal("server.SetZapDomain:", err)
}
err = server.Bind("tcp://127.0.0.1:9687")
if err != nil {
t.Fatal("server.Bind:", err)
}
err = client.Connect("tcp://127.0.0.1:9687")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
err = server.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("server.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
server.Unbind("tcp://127.0.0.1:9687")
client.Disconnect("tcp://127.0.0.1:9687")
// Now use the right domain, the test must pass
err = server.SetZapDomain("TEST")
if err != nil {
t.Fatal("server.SetZapDomain:", err)
}
err = server.Bind("tcp://127.0.0.1:9688")
if err != nil {
t.Fatal("server.Bind:", err)
}
err = client.Connect("tcp://127.0.0.1:9688")
if err != nil {
t.Fatal("client.Connect:", err)
}
msg, err = bounce(server, client)
if err != nil {
t.Error(msg, err)
}
server.Unbind("tcp://127.0.0.1:9688")
client.Disconnect("tcp://127.0.0.1:9688")
err = client.Close()
client = nil
if err != nil {
t.Error("client.Close:", err)
}
err = server.Close()
server = nil
if err != nil {
t.Error("server.Close:", err)
}
}
func TestSecurityCurve(t *testing.T) {
time.Sleep(100 * time.Millisecond)
var handler, server, client *zmq.Socket
defer func() {
for _, s := range []*zmq.Socket{handler} {
if s != nil {
s.SetLinger(0)
s.Close()
}
}
}()
if _, minor, _ := zmq.Version(); minor >= 1 && !zmq.HasCurve() {
t.Skip("Curve not available")
}
// Generate new keypairs for this test
client_public, client_secret, err := zmq.NewCurveKeypair()
if err != nil {
t.Fatal("NewCurveKeypair:", err)
}
server_public, server_secret, err := zmq.NewCurveKeypair()
if err != nil {
t.Fatal("NewCurveKeypair:", err)
}
handler, err = zmq.NewSocket(zmq.REP)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = handler.Bind("inproc://zeromq.zap.01")
if err != nil {
t.Fatal("handler.Bind:", err)
}
doHandler := func(state zmq.State) error {
msg, err := handler.RecvMessage(0)
if err != nil {
return err // Terminating
}
version := msg[0]
sequence := msg[1]
// domain := msg[2]
// address := msg[3]
identity := msg[4]
mechanism := msg[5]
client_key := msg[6]
client_key_text := zmq.Z85encode(client_key)
if version != "1.0" {
return errors.New("version != 1.0")
}
if mechanism != "CURVE" {
return errors.New("mechanism != CURVE")
}
if identity != "IDENT" {
return errors.New("identity != IDENT")
}
if client_key_text == client_public {
handler.SendMessage(version, sequence, "200", "OK", "anonymous", "")
} else {
handler.SendMessage(version, sequence, "400", "Invalid client public key", "", "")
}
return nil
}
doQuit := func(i interface{}) error {
err := handler.Close()
handler = nil
if err != nil {
t.Error("handler.Close:", err)
}
return errors.New("Quit")
}
quit := make(chan interface{})
reactor := zmq.NewReactor()
reactor.AddSocket(handler, zmq.POLLIN, doHandler)
reactor.AddChannel(quit, 0, doQuit)
go func() {
reactor.Run(100 * time.Millisecond)
quit <- true
}()
defer func() {
quit <- true
<-quit
close(quit)
}()
// Server socket will accept connections
server, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = server.SetCurveServer(1)
if err != nil {
t.Fatal("server.SetCurveServer(1):", err)
}
err = server.SetCurveSecretkey(server_secret)
if err != nil {
t.Fatal("server.SetCurveSecretkey:", err)
}
err = server.SetIdentity("IDENT")
if err != nil {
t.Fatal("server.SetIdentity:", err)
}
server.Bind("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("server.Bind:", err)
}
err = server.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("server.SetRcvtimeo:", err)
}
// Check CURVE security with valid credentials
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetCurveServerkey(server_public)
if err != nil {
t.Fatal("client.SetCurveServerkey:", err)
}
err = client.SetCurvePublickey(client_public)
if err != nil {
t.Fatal("client.SetCurvePublickey:", err)
}
err = client.SetCurveSecretkey(client_secret)
if err != nil {
t.Fatal("client.SetCurveSecretkey:", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
msg, err := bounce(server, client)
if err != nil {
t.Error(msg, err)
}
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
time.Sleep(100 * time.Millisecond)
// Check CURVE security with a garbage server key
// This will be caught by the curve_server class, not passed to ZAP
garbage_key := "0000111122223333444455556666777788889999"
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetCurveServerkey(garbage_key)
if err != nil {
t.Fatal("client.SetCurveServerkey:", err)
}
err = client.SetCurvePublickey(client_public)
if err != nil {
t.Fatal("client.SetCurvePublickey:", err)
}
err = client.SetCurveSecretkey(client_secret)
if err != nil {
t.Fatal("client.SetCurveSecretkey:", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
client.SetLinger(0)
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
time.Sleep(100 * time.Millisecond)
// Check CURVE security with a garbage client secret key
// This will be caught by the curve_server class, not passed to ZAP
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetCurveServerkey(server_public)
if err != nil {
t.Fatal("client.SetCurveServerkey:", err)
}
err = client.SetCurvePublickey(garbage_key)
if err != nil {
t.Fatal("client.SetCurvePublickey:", err)
}
err = client.SetCurveSecretkey(client_secret)
if err != nil {
t.Fatal("client.SetCurveSecretkey:", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
client.SetLinger(0)
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
time.Sleep(100 * time.Millisecond)
// Check CURVE security with a garbage client secret key
// This will be caught by the curve_server class, not passed to ZAP
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetCurveServerkey(server_public)
if err != nil {
t.Fatal("client.SetCurveServerkey:", err)
}
err = client.SetCurvePublickey(client_public)
if err != nil {
t.Fatal("client.SetCurvePublickey:", err)
}
err = client.SetCurveSecretkey(garbage_key)
if err != nil {
t.Fatal("client.SetCurveSecretkey:", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
client.SetLinger(0)
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
time.Sleep(100 * time.Millisecond)
// Check CURVE security with bogus client credentials
// This must be caught by the ZAP handler
bogus_public, bogus_secret, _ := zmq.NewCurveKeypair()
client, err = zmq.NewSocket(zmq.DEALER)
if err != nil {
t.Fatal("NewSocket:", err)
}
err = client.SetCurveServerkey(server_public)
if err != nil {
t.Fatal("client.SetCurveServerkey:", err)
}
err = client.SetCurvePublickey(bogus_public)
if err != nil {
t.Fatal("client.SetCurvePublickey:", err)
}
err = client.SetCurveSecretkey(bogus_secret)
if err != nil {
t.Fatal("client.SetCurveSecretkey:", err)
}
err = client.Connect("tcp://127.0.0.1:9998")
if err != nil {
t.Fatal("client.Connect:", err)
}
err = client.SetRcvtimeo(time.Second)
if err != nil {
t.Fatal("client.SetRcvtimeo:", err)
}
_, err = bounce(server, client)
if err == nil {
t.Error("Expected failure, got success")
}
client.SetLinger(0)
err = client.Close()
client = nil
if err != nil {
t.Fatal("client.Close:", err)
}
// Shutdown
err = server.Close()
server = nil
if err != nil {
t.Error("server.Close:", err)
}
}
func (w *Worker) run(socket *zmq.Socket) {
defer close(w.wait)
defer socket.Close()
socket.SetRcvtimeo(w.passiveTimeout)
socket.Bind(w.address)
atomic.StoreInt32(&w.runningWorkers, 0)
responseChannel := make(chan *messageContext)
sendResponse := func(response [][]byte) {
atomic.AddInt32(&w.runningWorkers, -1)
if _, err := socket.SendMessage(response); err != nil {
log.Println("Failed to send response:", err)
}
if atomic.LoadInt32(&w.runningWorkers) == 0 {
socket.SetRcvtimeo(w.passiveTimeout)
}
}
for {
if atomic.LoadInt32(&w.runningWorkers) == w.maxWorkers {
// We're already running maxWorkers so block until a response is ready
select {
case response := <-responseChannel:
w.logFinish(response)
sendResponse(response.data)
case <-w.quit:
return
}
}
select {
case <-w.quit:
return
case response := <-responseChannel:
w.logFinish(response)
sendResponse(response.data)
break
default:
message, err := socket.RecvMessageBytes(0)
if err != nil {
// Needed to yield to goroutines when GOMAXPROCS is 1.
// Note: The 1.3 preemptive scheduler doesn't seem to work here,
// so this is still required.
runtime.Gosched()
break
}
request, err := w.unmarshal(message[len(message)-1])
if err != nil {
log.Println("Received invalid message", err)
break
}
workerFunction := w.registeredWorkerFunctions[request.Method]
if workerFunction == nil {
log.Println("Unregistered worker function:", request.Method)
break
}
context := messageContext{data: message, startTime: time.Now(), method: request.Method}
if atomic.LoadInt32(&w.runningWorkers) == 0 {
socket.SetRcvtimeo(w.activeTimeout)
}
atomic.AddInt32(&w.runningWorkers, 1)
go w.runFunction(responseChannel, &context, request.Parameters, workerFunction)
}
}
}
func (server *server_t) ping(socket *zmq.Socket) {
if time.Now().After(server.ping_at) {
socket.SendMessage(server.endpoint, "PING")
server.ping_at = time.Now().Add(PING_INTERVAL)
}
}
func (re *RouterElement) run(doRouting func(*RouterElement, *Peer, [][]byte) (*Peer, [][]byte, error), controlIn, controlOut *zmq.Socket) {
var wg sync.WaitGroup
controlChanInternal := make(chan *controlMessage, 1)
wg.Add(1)
// Start a goroutine to copy from the control channel to the internal
// control channel and wake up the control socket pair on each message.
// Quit the goroutine when the control channel itself is closed
go func() {
defer func() {
controlIn.SetLinger(0)
controlIn.Close()
wg.Done()
}()
msg := [][]byte{[]byte{0}}
for {
select {
case m, ok := <-re.controlChan:
if !ok {
close(controlChanInternal)
// send a message to the control panel to wake it up
controlIn.SendMessage(msg)
// ensure we get a message back before closing it
controlIn.RecvMessageBytes(0)
return
}
controlChanInternal <- m
controlIn.SendMessageDontwait(msg)
}
}
}()
defer func() {
wg.Wait()
controlOut.SetLinger(0)
controlOut.Close()
// Close() should already remove all the peers, so this is just belt & braces
for _, p := range re.peers {
p.sock.SetLinger(0)
p.sock.Close()
close(p.buffer)
}
re.wg.Done()
}()
for {
poller := zmq.NewPoller()
pollSource := make([]*Peer, len(re.peers), len(re.peers))
blockInput := false
i := 0
for _, p := range re.peers {
pollSource[i] = p
i++
if p.xmit == nil {
select {
case p.xmit = <-p.buffer:
default:
}
}
buffered := len(p.buffer)
if p.xmit != nil {
buffered++
}
if buffered >= cap(p.buffer) {
log.Printf("Blocking input as %s cannot take more data", p.Name)
blockInput = true
}
}
for _, p := range pollSource {
flags := zmq.State(0)
if !blockInput {
flags |= zmq.POLLIN
}
if p.xmit != nil {
flags |= zmq.POLLOUT
}
poller.Add(p.sock, flags)
}
idxControl := poller.Add(controlOut, zmq.POLLIN)
if polled, err := poller.PollAll(-1); err != nil {
log.Printf("Cannot poll: %v", err)
} else {
for i, p := range pollSource {
if (polled[i].Events&zmq.POLLOUT != 0) && (p.xmit != nil) {
if _, err := p.sock.SendMessageDontwait(p.xmit); err != nil {
log.Printf("Peer %s cannot send message: %v", p.Name, err)
}
p.xmit = nil
}
}
for i, p := range pollSource {
if polled[i].Events&zmq.POLLIN != 0 {
if recv, err := p.sock.RecvMessageBytes(0); err != nil {
log.Printf("Peer %s cannot receive message: %v", p.Name, err)
} else {
if dest, fwd, err := doRouting(re, p, recv); err != nil {
log.Printf("Peer %s cannot route message: %v", p.Name, err)
} else {
if dest != nil && fwd != nil {
dest.buffer <- fwd
}
}
}
}
}
if polled[idxControl].Events&zmq.POLLIN != 0 {
// read from the control socket and discard. We don't care
// if this errors as it's really only to make us read from our
// own control socket
controlOut.RecvMessageBytes(0)
select {
case msg, ok := <-controlChanInternal:
if !ok {
// control channel is closed
// send a dummy message to handshake
msg := [][]byte{[]byte{0}}
controlOut.SendMessage(msg)
return
}
msg.reply <- re.processControlMessage(msg)
default:
}
}
}
}
}
func TestRouterElement(t *testing.T) {
var c *zmq.Context
var m *Mirror
var sa, sb *zmq.Socket
var err error
var addra string
var addrb string
defer func() {
if sa != nil {
sa.Close()
}
if sb != nil {
sb.Close()
}
if m != nil {
m.Close()
}
if c != nil {
c.Term()
}
}()
if c, err = zmq.NewContext(); err != nil {
t.Fatalf("Failed to create ZMQ context: %v", err)
}
if m, err = NewMirror(t, c, 1); err != nil {
t.Fatalf("Failed to create a new mirror: %v", err)
}
if sa, addra, err = svcrouter.NewHalfPair(c, true); err != nil {
t.Fatalf("Failed to create half pair A: %v", err)
}
if sb, addrb, err = svcrouter.NewHalfPair(c, true); err != nil {
t.Fatalf("Failed to create half pair B: %v", err)
}
if err = m.AddPeer(svcrouter.PeerDefinition{
Name: "ab",
ZmqType: zmq.PAIR,
Address: addra,
Bind: false,
PeerImpl: &MirrorPeerImpl{},
}); err != nil {
t.Fatalf("Could not add peer A: %v", err)
}
if err = m.AddPeer(svcrouter.PeerDefinition{
Name: "ba",
ZmqType: zmq.PAIR,
Address: addrb,
Bind: false,
PeerImpl: &MirrorPeerImpl{},
}); err != nil {
t.Fatalf("Could not add peer B: %v", err)
}
svcrouter.Barrier()
num := 100
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < num; i++ {
if _, err := sb.RecvMessage(0); err != nil {
t.Fatalf("sb receive error: %v", err)
}
}
}()
msg := [][]byte{[]byte("Hello"), []byte("World")}
for i := 0; i < num; i++ {
if _, err := sa.SendMessage(msg); err != nil {
t.Fatalf("sa send error: %v", err)
}
}
wg.Wait()
}
// NewWriterSocketHandler returns a function suitable for use as a handler
// by zmq.Reactor
func NewWriterSocketHandler(writerSocket *zmq4.Socket,
messageChan chan<- types.Message) func(zmq4.State) error {
// these messages get only and ack, not a reply
var ackOnlyMessages = map[string]ackOnlyMessageHandler{
"ping": handlePing,
"resilient-server-handshake": handleHandshake,
"resilient-server-signoff": handleSignoff}
return func(_ zmq4.State) error {
var err error
var ok bool
var rawMessage []string
var message types.Message
if rawMessage, err = writerSocket.RecvMessage(0); err != nil {
return fmt.Errorf("RecvMessage %s", err)
}
message.Marshalled = rawMessage[0]
message.Data = []byte(strings.Join(rawMessage[1:], ""))
if message.Type, err = msg.GetMessageType(message.Marshalled); err != nil {
return fmt.Errorf("GetMessageType %s", err)
}
if message.ID, err = msg.GetMessageID(message.Marshalled); err != nil {
return fmt.Errorf("GetMessageID %s", err)
}
reply := map[string]interface{}{
"message-type": "resilient-server-ack",
"message-id": message.ID,
"incoming-type": message.Type,
"accepted": true}
marshalledReply, err := json.Marshal(reply)
if err != nil {
return fmt.Errorf("Marshal %s", err)
}
_, err = writerSocket.SendMessage([]string{string(marshalledReply)})
if err != nil {
return fmt.Errorf("SendMessage %s", err)
}
/* ---------------------------------------------------------------
** 2014-05-21 dougfort:
** The python version of ResiliantServer maintains a dict of
** "client-address" for sending replies, but IMO we don't need
** that here because every message contains "client-address"
** so we can decouple the reply.
** -------------------------------------------------------------*/
handler, ok := ackOnlyMessages[message.Type]
if ok {
handler(message)
return nil
}
fog.Debug("writer-socket-handler received %s %s",
message.Type, message.ID)
messageChan <- message
return nil
}
}
