MB-61292: Handle encryption state on file open

Anytime logger opens a new file, initialize encryption state for
the disk sink. Add an encryption header if disk sink should be
encrypted. Also handle existing files and append cases when
log files are opened for continuation.

Change-Id: I08aa6bef2284575c895e013428f3b2d811537307
Reviewed-on: https://review.couchbase.org/c/ns_server/+/216152
Tested-by: Navdeep S Boparai <[email protected]>
Well-Formed: Build Bot <[email protected]>
Reviewed-by: Timofey Barmin <[email protected]>

Author: boparai11

apps/ale/src/ale.erl

@@ -29,7 +29,10 @@
          %% Callbacks for encryption
          create_no_deks_snapshot/0,
          file_encrypt_state_match/2,
+         is_file_encr_by_ds/2,
+         is_file_encrypted/1,
          file_encrypt_init/2,
+         file_encrypt_cont/3,
          file_encrypt_chunk/2,
 
          with_configuration_batching/1,
@@ -211,10 +214,22 @@ file_encrypt_state_match(DS, EncrState) ->
     Callback = get_encryption_cb(file_encrypt_state_match),
     Callback(DS, EncrState).
 
+is_file_encr_by_ds(Path, DS) ->
+    Callback = get_encryption_cb(is_file_encr_by_ds),
+    Callback(Path, DS).
+
+is_file_encrypted(Path) ->
+    Callback = get_encryption_cb(is_file_encrypted),
+    Callback(Path).
+
 file_encrypt_init(FileName, DS) ->
     Callback = get_encryption_cb(file_encrypt_init),
     Callback(FileName, DS).
 
+file_encrypt_cont(FileName, Offset, DS) ->
+    Callback = get_encryption_cb(file_encrypt_cont),
+    Callback(FileName, Offset, DS).
+
 file_encrypt_chunk(Data, EncrState) ->
     Callback = get_encryption_cb(file_encrypt_chunk),
     Callback(Data, EncrState).
@@ -431,10 +446,22 @@ default_encr_disabled_cbs() ->
           fun(_DS, _EncrState) ->
                   true
           end,
+      is_file_encr_by_ds =>
+          fun(_Path, _DS) ->
+                  true
+          end,
+      is_file_encrypted =>
+          fun(_Path) ->
+                  false
+          end,
       file_encrypt_init =>
           fun(_FileName, _DS) ->
                   {<<>>, #{}}
           end,
+      file_encrypt_cont =>
+          fun(_FileName, _Offset, _DS) ->
+                  #{}
+          end,
       file_encrypt_chunk =>
           fun(Data, EncrState) ->
                   {Data, EncrState}

apps/ale/src/ale_disk_sink.erl

@@ -308,20 +308,57 @@ maybe_rotate_files(#worker_state{file_size = FileSize,
 maybe_rotate_files(State) ->
     State.
 
+maybe_write_header(_SinkName, _File, <<>>) ->
+    ok;
+maybe_write_header(SinkName, File, Header) ->
+    time_stat(SinkName, write_time,
+              fun () ->
+                      ok = file:write(File, Header)
+              end).
+
+update_file_encr_state(true = _ActiveKeyMatch,
+                       #worker_state{path = Path} = State, DS) ->
+    {ok, File, #file_info{size = Size, inode = Inode}} = open_file(Path),
+    ContEncrState =
+        ale:file_encrypt_cont(filename:basename(Path), Size, DS),
+    State#worker_state{file = File,
+                       file_size = Size,
+                       file_inode = Inode,
+                       encr_state = ContEncrState};
+update_file_encr_state(false = _ActiveKeyMatch, State, _DS) ->
+    do_rotate_files(State).
+
+open_with_encr_state(true = _IsNewFile,
+                     #worker_state{sink_name = SinkName,
+                                   path = Path} = State, DS) ->
+    {ok, File, #file_info{size = 0, inode = Inode}} = open_file(Path),
+    {Header, EncrState} =
+        ale:file_encrypt_init(filename:basename(Path), DS),
+    maybe_write_header(SinkName, File, Header),
+    State#worker_state{file = File,
+                       file_size = byte_size(Header),
+                       file_inode = Inode,
+                       encr_state = EncrState};
+open_with_encr_state(false = _IsNewFile,
+                     #worker_state{path = Path} = State, DS) ->
+    ActiveKeyMatch = ale:is_file_encr_by_ds(Path, DS),
+    update_file_encr_state(ActiveKeyMatch, State, DS).
+
 open_log_file(#worker_state{path = Path,
-                            file = OldFile} = State) ->
+                            file = OldFile,
+                            sink_name = SinkName} = State) ->
     case OldFile of
         undefined ->
             ok;
         _ ->
             file:close(OldFile)
     end,
 
-    {ok, File, #file_info{size = Size,
-                          inode = Inode}} = open_file(Path),
-    State#worker_state{file = File,
-                       file_size = Size,
-                       file_inode = Inode}.
+    IsNewFile = not filelib:is_file(Path) orelse
+                filelib:file_size(Path) =:= 0,
+
+    DS = ale:get_sink_ds(SinkName),
+    open_with_encr_state(IsNewFile, State#worker_state{file = undefined}, DS).
 
 check_log_file(#worker_state{path = Path,
                              file_inode = FileInode} = State) ->
@@ -436,19 +473,17 @@ spawn_worker(WorkerState) ->
               worker_init(WorkerState)
       end).
 
-worker_init(#worker_state{rotation_check_interval = RotCheckInterval,
-                          path = Path} = State0) ->
+worker_init(#worker_state{
+               rotation_check_interval = RotCheckInterval} = State0) ->
     case RotCheckInterval > 0 of
         true ->
             erlang:send_after(RotCheckInterval, self(), check_file);
         false ->
             ok
     end,
 
-    DS = ale:create_no_deks_snapshot(),
-    {<<>>, EncrState} = ale:file_encrypt_init(filename:basename(Path), DS),
     worker_loop(
-        open_log_file(State0#worker_state{encr_state = EncrState})).
+      open_log_file(State0)).
 
 worker_loop(#worker_state{sink_name = SinkName} = State) ->
     NewState =
@@ -467,7 +502,7 @@ worker_loop(#worker_state{sink_name = SinkName} = State) ->
                 #worker_state{encr_state = EncrState} = State,
                 DS = ale:get_sink_ds(SinkName),
                 UpdtReq = not ale:file_encrypt_state_match(DS, EncrState),
-                UpdatedState = process_key_update_work(UpdtReq, DS, State),
+                UpdatedState = process_key_update_work(UpdtReq, State),
                 gen_server:reply(From, ok),
                 UpdatedState;
             Msg ->
@@ -508,23 +543,10 @@ write_data(InputData, InputDataSize,
                                   encr_state = NewEncrState},
     maybe_rotate_files(NewState).
 
-process_key_update_work(false = _UpdtReq, _DS, State) ->
+process_key_update_work(false = _UpdtReq, State) ->
     State;
-process_key_update_work(true = _UpdtReq, DS,
-                        #worker_state{sink_name = Name,
-                                      path = Path} = State) ->
-    {Header, EncryptState} =
-        ale:file_encrypt_init(filename:basename(Path), DS),
-    NewState = do_rotate_files(State),
-    #worker_state{file = File,
-                  file_size = 0,
-                  path = Path} = NewState,
-    time_stat(Name, write_time,
-              fun () ->
-                      ok = file:write(File, Header)
-              end),
-    NewState#worker_state{file_size = byte_size(Header),
-                          encr_state = EncryptState}.
+process_key_update_work(true = _UpdtReq, State) ->
+    do_rotate_files(State).
 
 remove_unnecessary_log_files(LogFilePath, NumFiles) ->
     Dir = filename:dirname(LogFilePath),

apps/ns_server/include/ns_common.hrl

@@ -470,10 +470,22 @@
           fun(DS, EncrState) ->
                   cb_crypto:file_encrypt_state_match(DS, EncrState)
           end,
+      is_file_encr_by_ds =>
+          fun(Path, DS) ->
+                  cb_crypto:is_file_encr_by_ds(Path, DS)
+          end,
+      is_file_encrypted =>
+          fun(Path) ->
+                  cb_crypto:is_file_encrypted(Path)
+          end,
       file_encrypt_init =>
           fun(FileName, DS) ->
                   cb_crypto:file_encrypt_init(FileName, DS)
           end,
+      file_encrypt_cont =>
+          fun(FileName, Offset, DS) ->
+                  cb_crypto:file_encrypt_cont(FileName, Offset, DS)
+          end,
       file_encrypt_chunk =>
           fun(Data, EncrState) ->
                   cb_crypto:file_encrypt_chunk(Data, EncrState)

apps/ns_server/src/cb_crypto.erl

@@ -29,10 +29,13 @@
          atomic_write_file/3,
          atomic_write_file/4,
          file_encrypt_init/2,
+         file_encrypt_cont/3,
          file_encrypt_chunk/2,
 
          read_file/2,
          file_encrypt_state_match/2,
+         is_file_encr_by_ds/2,
+         is_file_encrypted/1,
          file_decrypt_init/3,
          file_decrypt_next_chunk/2,
 
@@ -114,18 +117,6 @@ atomic_write_file(Path, Bytes, DekSnapshot, Opts) when is_binary(Bytes) ->
           encrypt_to_file(IODevice, Filename, Bytes, MaxChunkSize, DekSnapshot)
       end).
 
-get_key_id(undefined) ->
-    undefined;
-get_key_id(#{id := KeyId} = _Key) ->
-    KeyId.
-
--spec file_encrypt_state_match(#dek_snapshot{},
-                               #file_encr_state{}) -> boolean().
-file_encrypt_state_match(#dek_snapshot{active_key = ActiveKey},
-                         #file_encr_state{key = FileActiveKey}) ->
-
-    get_key_id(ActiveKey) =:= get_key_id(FileActiveKey).
-
 -spec file_encrypt_init(string(), #dek_snapshot{}) ->
           {binary(), #file_encr_state{}}.
 file_encrypt_init(Filename,
@@ -149,6 +140,18 @@ file_encrypt_init(Filename,
                               iv_atomic_counter = IVCounter,
                               offset = size(Header)}}.
 
+-spec file_encrypt_cont(string(), non_neg_integer(), #dek_snapshot{}) ->
+          #file_encr_state{}.
+file_encrypt_cont(Filename, Offset,
+                  #dek_snapshot{active_key = ActiveKey,
+                                iv_random = IVRandom,
+                                iv_atomic_counter = IVCounter}) ->
+    #file_encr_state{filename = iolist_to_binary(Filename),
+                     key = ActiveKey,
+                     iv_random = IVRandom,
+                     iv_atomic_counter = IVCounter,
+                     offset = Offset}.
+
 -spec file_encrypt_chunk(binary(), #file_encr_state{}) ->
           {binary(), #file_encr_state{}}.
 file_encrypt_chunk(Data, #file_encr_state{key = undefined} = State) ->
@@ -164,6 +167,32 @@ file_encrypt_chunk(Data, #file_encr_state{key = Dek,
     NewOffset = Offset + size(ChunkWithSize),
     {ChunkWithSize, State#file_encr_state{offset = NewOffset}}.
 
+-spec file_encrypt_state_match(#dek_snapshot{},
+                               #file_encr_state{}) -> boolean().
+file_encrypt_state_match(#dek_snapshot{active_key = ActiveKey},
+                         #file_encr_state{key = FileActiveKey}) ->
+    get_key_id(ActiveKey) =:= get_key_id(FileActiveKey).
+
+-spec is_file_encr_by_ds(string(), #dek_snapshot{}) -> true | false.
+is_file_encr_by_ds(Path, #dek_snapshot{active_key = undefined}) ->
+    not is_file_encrypted(Path);
+is_file_encr_by_ds(Path, #dek_snapshot{active_key = #{id := KeyId}}) ->
+    {ok, File} = file:open(Path, [read, raw, binary]),
+    try
+        header_key_match(file:read(File, ?ENCRYPTED_FILE_HEADER_LEN), KeyId)
+    after
+        file:close(File)
+    end.
+
+-spec is_file_encrypted(string()) -> true | false.
+is_file_encrypted(Path) ->
+    {ok, File} = file:open(Path, [read, raw, binary]),
+    try
+        is_valid_encr_header(file:read(File, ?ENCRYPTED_FILE_HEADER_LEN))
+    after
+        file:close(File)
+    end.
+
 -spec read_file(string(), cb_deks:dek_kind() | fun(() -> fetch_deks_res())) ->
           {decrypted, binary()} | {raw, binary()} | {error, term()}.
 read_file(Path, GetDekSnapshotFun) when is_function(GetDekSnapshotFun, 0) ->
@@ -352,6 +381,31 @@ get_dek_rotation_interval(Type, Snapshot) ->
 %%% Internal functions
 %%%===================================================================
 
+is_valid_encr_header({ok, HeaderData}) ->
+    case parse_header(HeaderData) of
+        {ok, _Parsed} ->
+            true;
+        _ ->
+            false
+    end;
+is_valid_encr_header(eof) ->
+    false.
+
+header_key_match({ok, HeaderData}, KeyId) ->
+    case parse_header(HeaderData) of
+        {ok, {0, KeyId, ?ENCRYPTED_FILE_HEADER_LEN, _Rest}} ->
+            true;
+        _ ->
+            false
+    end;
+header_key_match(eof, _KeyId) ->
+    false.
+
+get_key_id(undefined) ->
+    undefined;
+get_key_id(#{id := KeyId} = _Key) ->
+    KeyId.
+
 encrypt_internal(Data, AD, IVRandom, IVAtomic, #{type := 'raw-aes-gcm',
                                                  info := #{key := KeyFun}}) ->
     IV = new_aes_gcm_iv(IVRandom, IVAtomic),