Golang util.MessageStream类代码示例

// RunMessageLoop handles messages from a client on a given message stream until
// it gets an error trying to read a message.
func (m *ResourceMaster) RunMessageLoop(ms util.MessageStream, programPolicy *ProgramPolicy) error {
	for {
		var msg Message
		if err := ms.ReadMessage(&msg); err != nil {
			return err
		}

		var fop FileOperation
		t := *msg.Type
		if t == MessageType_CREATE || t == MessageType_DELETE || t == MessageType_READ || t == MessageType_WRITE {
			if err := proto.Unmarshal(msg.Data, &fop); err != nil {
				log.Printf("Couldn't unmarshal FileOperation for operation %d\n", t)
				continue
			}

			if err := m.checkFileAuth(&msg, &fop); err != nil {
				log.Printf("The file operation %d didn't pass authorization: %s\n", t, err)
				continue
			}
		}

		switch *msg.Type {
		case MessageType_AUTH_CERT:
			cert, err := m.AuthenticatePrincipal(ms, &msg, programPolicy)
			if err != nil {
				log.Printf("Failed to authenticate a principal: %s\n", err)
				continue
			}

			owner, err := principalNameFromDERCert(cert)
			if err != nil {
				log.Printf("Couldn't get the owner name from the cert: %s\n", err)
				continue
			}
			_, err = m.InsertPrincipal(owner, cert, Authenticated)
			if err != nil {
				log.Printf("Couldn't set the principal as authenticated")
			}
		case MessageType_CREATE:
			if err := m.Create(ms, &fop); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		case MessageType_READ:
			if err := m.Read(ms, &fop, programPolicy.SymKeys); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		case MessageType_WRITE:
			if err := m.Write(ms, &fop, programPolicy.SymKeys); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		default:
			if err := sendResult(ms, false); err != nil {
				log.Printf("Couldn't signal failure for the invalid operation: %s", err)
			}
			log.Printf("Invalid initial message type %d\n", *msg.Type)
		}
	}
}

// checkResponse waits for a RollbackResponse and checks to make sure it's not
// an ERROR response from the server.
func checkResponse(ms *util.MessageStream) error {
	var rr RollbackResponse
	if err := ms.ReadMessage(&rr); err != nil {
		return err
	}
	if *rr.Type == RollbackMessageType_ERROR {
		return fmt.Errorf("couldn't set the counter on the remote server")
	}
	return nil
}

func sendError(err error, ms *util.MessageStream) {
	var errStr = ""
	if err != nil {
		errStr = err.Error()
	}
	resp := &domain_service.DomainServiceResponse{ErrorMessage: &errStr}
	if _, err := ms.WriteMessage(resp); err != nil {
		log.Printf("domain_server: Error sending resp on the channel: %s\n ", err)
	}
}

func SendMessage(ms *util.MessageStream, msg *SimpleMessage) error {
	out, err := proto.Marshal(msg)
	if err != nil {
		return errors.New("SendRequest: Can't encode response")
	}
	send := string(out)
	_, err = ms.WriteString(send)
	if err != nil {
		return errors.New("SendResponse: Writestring error")
	}
	return nil
}

// recvResult waits for a OperationResult on a MessageStream
func recvResult(ms util.MessageStream) (bool, error) {
	var m Message
	if err := ms.ReadMessage(&m); err != nil {
		return false, err
	}
	var res OperationResult
	if err := proto.Unmarshal(m.Data, &res); err != nil {
		return false, err
	}

	return *res.Result, nil
}

func GetMessage(ms *util.MessageStream) (*SimpleMessage,
	error) {
	resp, err := ms.ReadString()
	if err != nil {
		return nil, err
	}
	msg := new(SimpleMessage)
	err = proto.Unmarshal([]byte(resp), msg)
	if err != nil {
		return nil, errors.New("GetResponse: Can't unmarshal message")
	}
	return msg, nil
}

// SetCounter sets the remote counter for this program.
func SetCounter(ms *util.MessageStream, counter uint64) error {
	rm := &RollbackMessage{
		Type: RollbackMessageType_SET_COUNTER.Enum(),
		Data: EncodeCounter(counter),
	}
	if _, err := ms.WriteMessage(rm); err != nil {
		return err
	}

	// TODO(tmroeder): we currently ignore the value of the counter returned
	// by the server.
	return checkResponse(ms)
}

// readResult reads an OperationResult and returns its value or an error.
func readResult(ms util.MessageStream) (bool, error) {
	// Read the response wrapper message.
	var arm Message
	if err := ms.ReadMessage(&arm); err != nil {
		return false, err
	}
	if *arm.Type != MessageType_OP_RES {
		return false, fmt.Errorf("didn't receive OP_RES from the server")
	}

	var opr OperationResult
	if err := proto.Unmarshal(arm.Data, &opr); err != nil {
		return false, err
	}
	return *opr.Result, nil
}

// sendResult sends an OperationResult with the given value on the given stream.
func sendResult(ms util.MessageStream, result bool) error {
	res := &OperationResult{Result: proto.Bool(result)}
	ar := &Message{
		Type: MessageType_OP_RES.Enum(),
	}
	var err error
	if ar.Data, err = proto.Marshal(res); err != nil {
		return err
	}

	if _, err := ms.WriteMessage(ar); err != nil {
		return err
	}

	return nil
}

func (conn *Conn) sendCredentials(ms util.MessageStream) {
	if conn.Authenticate && conn.delegation != nil {
		ms.WriteString("delegation")
		ms.WriteMessage(conn.delegation)
	} else if conn.Authenticate {
		ms.WriteString("key")
	} else {
		ms.WriteString("anonymous")
	}
}

// RunMessageLoop handles incoming messages for the RollbackMaster and passes
// them to the appropriate functions.
func (r *RollbackMaster) RunMessageLoop(ms *util.MessageStream, programPolicy *ProgramPolicy, name string) error {
	for {
		var msg RollbackMessage
		if err := ms.ReadMessage(&msg); err != nil {
			return err
		}

		switch *msg.Type {
		case RollbackMessageType_SET_COUNTER:
			i, err := decodeCounter(msg.Data)
			if err != nil {
				log.Printf("failed to decode counter for SET_COUNTER: %s", err)
				continue
			}

			if err = r.SetCounter(ms, name, i); err != nil {
				log.Printf("failed to set the counter on the RollbackMaster: %s", err)
				continue
			}
		case RollbackMessageType_GET_COUNTER:
			if err := r.GetCounter(ms, name); err != nil {
				log.Printf("failed to get the counter for program %s", name)
				continue
			}
		case RollbackMessageType_SET_HASH:
			var rh RollbackHash
			if err := proto.Unmarshal(msg.Data, &rh); err != nil {
				log.Printf("failed to unmarshal the parameters for SET_HASH: %s", err)
				continue
			}
			if err := r.SetHash(ms, name, *rh.Item, rh.Hash); err != nil {
				log.Printf("failed to set the hash for item %s on program %s: %s", *rh.Item, name, err)
				continue
			}
		case RollbackMessageType_GET_HASHED_VERIFIER:
			if err := r.GetHashedVerifier(ms, name, string(msg.Data)); err != nil {
				log.Printf("failed to get the hashed verifier for program %s: %s", name, err)
				continue
			}
		default:
			log.Printf("unknown rollback message %d", *msg.Type)
		}
	}
}

// SetHash sets the value of a hash for a given item for this program.
func SetHash(ms *util.MessageStream, item string, hash []byte) error {
	rh := &RollbackHash{
		Item: proto.String(item),
		Hash: hash,
	}
	rhb, err := proto.Marshal(rh)
	if err != nil {
		return err
	}
	rm := &RollbackMessage{
		Type: RollbackMessageType_SET_HASH.Enum(),
		Data: rhb,
	}
	if _, err := ms.WriteMessage(rm); err != nil {
		return err
	}

	return checkResponse(ms)
}

// GetCounter implements RollbackMessageType_GET_COUNTER and returns the current
// value of a counter to the requestor.
func (r *RollbackMaster) GetCounter(ms *util.MessageStream, name string) error {
	var emptyData []byte
	rr := &RollbackResponse{
		Type: RollbackMessageType_ERROR.Enum(),
		Data: emptyData,
	}

	p := r.FindRollbackProgram(name)
	if p == nil {
		if _, err := ms.WriteMessage(rr); err != nil {
			return err
		}
		return fmt.Errorf("couldn't find a rollback program with name %s", name)
	}

	rr.Type = RollbackMessageType_GET_COUNTER.Enum()
	rr.Data = EncodeCounter(p.MonotonicCounter)
	_, err := ms.WriteMessage(rr)
	return err
}

// GetHashedVerifier gets the hash of the counter and the item hash for a given
// item.
func GetHashedVerifier(ms *util.MessageStream, item string) ([]byte, error) {
	rm := &RollbackMessage{
		Type: RollbackMessageType_GET_HASHED_VERIFIER.Enum(),
		Data: []byte(item),
	}
	if _, err := ms.WriteMessage(rm); err != nil {
		return nil, err
	}

	// We can't use checkResponse here since we need to get the value out of
	// the response to read the hash.
	var rr RollbackResponse
	if err := ms.ReadMessage(&rr); err != nil {
		return nil, err
	}
	if *rr.Type == RollbackMessageType_ERROR {
		return nil, fmt.Errorf("couldn't set the counter on the remote server")
	}

	return rr.Data, nil
}

// sendOperation is a helper method that sets up the data structures needed for
// a FileOperation message like CREATE, WRITE, or READ, and sends this message
// on the MessageStream.
func sendOperation(ms util.MessageStream, mt MessageType, cert []byte, name string) error {
	fop := &FileOperation{
		Subject: cert,
		Name:    proto.String(name),
	}

	fopb, err := proto.Marshal(fop)
	if err != nil {
		return err
	}
	m := &Message{
		Type: mt.Enum(),
		Data: fopb,
	}

	if _, err := ms.WriteMessage(m); err != nil {
		return err
	}

	return nil
}