From 75d7e656727da7422320c1ead4a73e996be09f3c Mon Sep 17 00:00:00 2001
From: Ery Lee <ery.lee@gmail.com>
Date: Fri, 16 Jan 2015 12:35:12 +0800
Subject: [PATCH] set_flag, unset_flag

---
 apps/emqtt/src/emqtt_message.erl | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/apps/emqtt/src/emqtt_message.erl b/apps/emqtt/src/emqtt_message.erl
index 3ecbd54da..aa4c415f0 100644
--- a/apps/emqtt/src/emqtt_message.erl
+++ b/apps/emqtt/src/emqtt_message.erl
@@ -30,6 +30,8 @@
 
 -export([from_packet/1, to_packet/1]).
 
+-export([set_flag/1, set_flag/2, unset_flag/1, unset_flag/2]).
+
 %%----------------------------------------------------------------------------
 
 -ifdef(use_specs).
@@ -38,6 +40,10 @@
 
 -spec( to_packet( mqtt_message() ) -> mqtt_packet() ).
 
+-sepc( set_flag(atom(), mqtt_message() ) -> mqtt_message().
+
+-sepc( unset_flag(atom(), mqtt_message() ) -> mqtt_message().
+
 -endif.
 
 %%----------------------------------------------------------------------------
@@ -95,3 +101,22 @@ to_packet(#mqtt_message{ msgid   = MsgId,
                                                     packet_id  = PacketId }, 
                   payload = Payload }.
 
+%%
+%% @doc set dup, retain flag
+%%
+set_flag(Msg) ->
+    Msg#mqtt_message{dup = true, retain = true}.
+set_flag(dup, Msg = #mqtt_message{dup = false}) -> 
+    Msg#mqtt_message{dup = true};
+set_flag(retain, Msg = #mqtt_message{retain = false}) ->
+    Msg#mqtt_message{retain = true};
+set_flag(Flag, Msg) when Flag =:= dup orelse Flag =:= retain -> Msg.
+
+unset_flag(Msg) ->
+    Msg#mqtt_message{dup = false, retain = false}.
+unset_flag(dup, Msg = #mqtt_message{dup = true}) -> 
+    Msg#mqtt_message{dup = false};
+unset_flag(retain, Msg = #mqtt_message{retain = true}) ->
+    Msg#mqtt_message{retain = false};
+unset_flag(Flag, Msg) when Flag =:= dup orelse Flag =:= retain -> Msg.
+