github.com/yunabe/lgo@v0.0.0-20190709125917-42c42d410fdf/jupyter/gojupyterscaffold/message.go (about)

     1  package gojupyterscaffold
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/hmac"
     6  	"crypto/rand"
     7  	"crypto/sha256"
     8  	"encoding/hex"
     9  	"encoding/json"
    10  	"errors"
    11  	"fmt"
    12  
    13  	zmq "github.com/pebbe/zmq4"
    14  )
    15  
    16  var (
    17  	identityDelim = []byte("<IDS|MSG>")
    18  )
    19  
    20  type message struct {
    21  	Identity     [][]byte
    22  	Header       messageHeader
    23  	ParentHeader messageHeader
    24  	Metadata     interface{}
    25  	Content      interface{}
    26  }
    27  
    28  // http://jupyter-client.readthedocs.io/en/latest/messaging.html#general-message-format
    29  type messageHeader struct {
    30  	MsgID    string `json:"msg_id"`
    31  	Username string `json:"username"`
    32  	Session  string `json:"session"`
    33  	Date     string `json:"date"`
    34  	MsgType  string `json:"msg_type"`
    35  	Version  string `json:"version"`
    36  }
    37  
    38  func newContentForMsgType(header *messageHeader) interface{} {
    39  	switch header.MsgType {
    40  	case "execute_request":
    41  		return &ExecuteRequest{}
    42  	case "complete_request":
    43  		return &CompleteRequest{}
    44  	case "inspect_request":
    45  		return &InspectRequest{}
    46  	case "is_complete_request":
    47  		return &IsCompleteRequest{}
    48  	case "gofmt_request":
    49  		return &GoFmtRequest{}
    50  	}
    51  	return nil
    52  }
    53  
    54  func unmarshalJSONToInterface(b []byte, i interface{}) (interface{}, error) {
    55  	if i == nil {
    56  		tmp := make(map[string]interface{})
    57  		i = &tmp
    58  	}
    59  	if err := json.Unmarshal(b, i); err != nil {
    60  		return nil, err
    61  	}
    62  	return i, nil
    63  }
    64  
    65  func validateMessages(msgs [][]byte, key []byte) error {
    66  	if len(msgs) < 5 {
    67  		return fmt.Errorf("Too short messages: %d", len(msgs))
    68  	}
    69  	mac := hmac.New(sha256.New, key)
    70  	// header, parent header, metadata, and content are signed with hmac.
    71  	for _, msg := range msgs[1:5] {
    72  		mac.Write(msg)
    73  	}
    74  	// Decode the hex signature
    75  	sig := make([]byte, hex.DecodedLen(len(msgs[0])))
    76  	hex.Decode(sig, msgs[0])
    77  	// Verify the signature
    78  	if !hmac.Equal(mac.Sum(nil), sig) {
    79  		return errors.New("HMAC was invalid")
    80  	}
    81  	return nil
    82  }
    83  
    84  func (m *message) Unmarshal(bs [][]byte, key []byte) error {
    85  	delimIdx := -1
    86  	for i, b := range bs {
    87  		if bytes.Equal(b, identityDelim) {
    88  			delimIdx = i
    89  			break
    90  		}
    91  	}
    92  	if delimIdx < 0 {
    93  		return fmt.Errorf("Identity deliminator %s not found", identityDelim)
    94  	}
    95  	bodies := bs[delimIdx+1:]
    96  	if err := validateMessages(bodies, key); err != nil {
    97  		return err
    98  	}
    99  	m.Identity = bs[:delimIdx]
   100  	m.Header = messageHeader{}
   101  	if err := json.Unmarshal(bodies[1], &m.Header); err != nil {
   102  		return err
   103  	}
   104  	m.ParentHeader = messageHeader{}
   105  	if err := json.Unmarshal(bodies[2], &m.ParentHeader); err != nil {
   106  		return err
   107  	}
   108  	if i, err := unmarshalJSONToInterface(bodies[3], m.Metadata); err == nil {
   109  		m.Metadata = i
   110  	} else {
   111  		return err
   112  	}
   113  	m.Content = newContentForMsgType(&m.Header)
   114  	if i, err := unmarshalJSONToInterface(bodies[4], m.Content); err == nil {
   115  		m.Content = i
   116  	} else {
   117  		return err
   118  	}
   119  	return nil
   120  }
   121  
   122  func (m *message) Marshal(key []byte) (bs [][]byte, err error) {
   123  	bs = m.Identity
   124  	bs = append(bs, identityDelim)
   125  	chunks := []interface{}{&m.Header, &m.ParentHeader, m.Metadata, m.Content}
   126  	var bodies [][]byte
   127  	for _, chunk := range chunks {
   128  		var data []byte
   129  		if chunk == nil {
   130  			data = []byte("{}")
   131  		} else {
   132  			data, err = json.Marshal(chunk)
   133  			if err != nil {
   134  				return nil, err
   135  			}
   136  		}
   137  		bodies = append(bodies, data)
   138  	}
   139  	mac := hmac.New(sha256.New, key)
   140  	for _, body := range bodies {
   141  		mac.Write(body)
   142  	}
   143  	sig := mac.Sum(nil)
   144  	hexSig := make([]byte, hex.EncodedLen(len(sig)))
   145  	hex.Encode(hexSig, sig)
   146  
   147  	bs = append(bs, hexSig)
   148  	return append(bs, bodies...), nil
   149  }
   150  
   151  func (m *message) Send(sock *zmq.Socket, key []byte) error {
   152  	bs, err := m.Marshal(key)
   153  	if err != nil {
   154  		return fmt.Errorf("Failed to marshal kernelinfo: %v", err)
   155  	}
   156  	args := make([]interface{}, len(bs))
   157  	for i, b := range bs {
   158  		args[i] = b
   159  	}
   160  	if _, err := sock.SendMessage(args...); err != nil {
   161  		return err
   162  	}
   163  	return nil
   164  }
   165  
   166  func genMsgID() string {
   167  	var id [32]byte
   168  	if _, err := rand.Read(id[:]); err != nil {
   169  		panic(fmt.Sprintf("rand.Read failed: %v", err))
   170  	}
   171  	var hexID [64]byte
   172  	hex.Encode(hexID[:], id[:])
   173  	return string(hexID[:])
   174  }
   175  
   176  func newMessageWithParent(parent *message) *message {
   177  	// https://github.com/ipython/ipykernel/blob/master/ipykernel/kernelbase.py#L222
   178  	// idents should be copied from the parent.
   179  	msg := message{
   180  		Identity:     parent.Identity,
   181  		ParentHeader: parent.Header,
   182  	}
   183  	msg.Header.Session = parent.Header.Session
   184  	msg.Header.Version = "5.2"
   185  	msg.Header.Username = "username"
   186  	msg.Header.MsgID = genMsgID()
   187  	return &msg
   188  }