1%% @author Bob Ippolito <bob@mochimedia.com>
2%% @copyright 2010 Mochi Media, Inc.
3
4%% @doc Create temporary files and directories. Requires crypto to be started.
5
6-module(mochitemp).
7-export([gettempdir/0]).
8-export([mkdtemp/0, mkdtemp/3]).
9-export([rmtempdir/1]).
10%% -export([mkstemp/4]).
11-define(SAFE_CHARS, {$a, $b, $c, $d, $e, $f, $g, $h, $i, $j, $k, $l, $m,
12                     $n, $o, $p, $q, $r, $s, $t, $u, $v, $w, $x, $y, $z,
13                     $A, $B, $C, $D, $E, $F, $G, $H, $I, $J, $K, $L, $M,
14                     $N, $O, $P, $Q, $R, $S, $T, $U, $V, $W, $X, $Y, $Z,
15                     $0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $_}).
16-define(TMP_MAX, 10000).
17
18-include_lib("kernel/include/file.hrl").
19
20%% TODO: An ugly wrapper over the mktemp tool with open_port and sadness?
21%%       We can't implement this race-free in Erlang without the ability
22%%       to issue O_CREAT|O_EXCL. I suppose we could hack something with
23%%       mkdtemp, del_dir, open.
24%% mkstemp(Suffix, Prefix, Dir, Options) ->
25%%    ok.
26
27rmtempdir(Dir) ->
28    case file:del_dir(Dir) of
29        {error, eexist} ->
30            ok = rmtempdirfiles(Dir),
31            ok = file:del_dir(Dir);
32        ok ->
33            ok
34    end.
35
36rmtempdirfiles(Dir) ->
37    {ok, Files} = file:list_dir(Dir),
38    ok = rmtempdirfiles(Dir, Files).
39
40rmtempdirfiles(_Dir, []) ->
41    ok;
42rmtempdirfiles(Dir, [Basename | Rest]) ->
43    Path = filename:join([Dir, Basename]),
44    case filelib:is_dir(Path) of
45        true ->
46            ok = rmtempdir(Path);
47        false ->
48            ok = file:delete(Path)
49    end,
50    rmtempdirfiles(Dir, Rest).
51
52mkdtemp() ->
53    mkdtemp("", "tmp", gettempdir()).
54
55mkdtemp(Suffix, Prefix, Dir) ->
56    mkdtemp_n(rngpath_fun(Suffix, Prefix, Dir), ?TMP_MAX).
57
58
59
60mkdtemp_n(RngPath, 1) ->
61    make_dir(RngPath());
62mkdtemp_n(RngPath, N) ->
63    try make_dir(RngPath())
64    catch throw:{error, eexist} ->
65            mkdtemp_n(RngPath, N - 1)
66    end.
67
68make_dir(Path) ->
69    case file:make_dir(Path) of
70        ok ->
71            ok;
72        E={error, eexist} ->
73            throw(E)
74    end,
75    %% Small window for a race condition here because dir is created 777
76    ok = file:write_file_info(Path, #file_info{mode=8#0700}),
77    Path.
78
79rngpath_fun(Prefix, Suffix, Dir) ->
80    fun () ->
81            filename:join([Dir, Prefix ++ rngchars(6) ++ Suffix])
82    end.
83
84rngchars(0) ->
85    "";
86rngchars(N) ->
87    [rngchar() | rngchars(N - 1)].
88
89rngchar() ->
90    rngchar(crypto:rand_uniform(0, tuple_size(?SAFE_CHARS))).
91
92rngchar(C) ->
93    element(1 + C, ?SAFE_CHARS).
94
95%% @spec gettempdir() -> string()
96%% @doc Get a usable temporary directory using the first of these that is a directory:
97%%      $TMPDIR, $TMP, $TEMP, "/tmp", "/var/tmp", "/usr/tmp", ".".
98gettempdir() ->
99    gettempdir(gettempdir_checks(), fun normalize_dir/1).
100
101gettempdir_checks() ->
102    [{fun os:getenv/1, ["TMPDIR", "TMP", "TEMP"]},
103     {fun gettempdir_identity/1, ["/tmp", "/var/tmp", "/usr/tmp"]},
104     {fun gettempdir_cwd/1, [cwd]}].
105
106gettempdir_identity(L) ->
107    L.
108
109gettempdir_cwd(cwd) ->
110    {ok, L} = file:get_cwd(),
111    L.
112
113gettempdir([{_F, []} | RestF], Normalize) ->
114    gettempdir(RestF, Normalize);
115gettempdir([{F, [L | RestL]} | RestF], Normalize) ->
116    case Normalize(F(L)) of
117        false ->
118            gettempdir([{F, RestL} | RestF], Normalize);
119        Dir ->
120            Dir
121    end.
122
123normalize_dir(False) when False =:= false orelse False =:= "" ->
124    %% Erlang doesn't have an unsetenv, wtf.
125    false;
126normalize_dir(L) ->
127    Dir = filename:absname(L),
128    case filelib:is_dir(Dir) of
129        false ->
130            false;
131        true ->
132            Dir
133    end.
134
135%%
136%% Tests
137%%
138-ifdef(TEST).
139-include_lib("eunit/include/eunit.hrl").
140
141pushenv(L) ->
142    [{K, os:getenv(K)} || K <- L].
143popenv(L) ->
144    F = fun ({K, false}) ->
145                %% Erlang doesn't have an unsetenv, wtf.
146                os:putenv(K, "");
147            ({K, V}) ->
148                os:putenv(K, V)
149        end,
150    lists:foreach(F, L).
151
152gettempdir_fallback_test() ->
153    ?assertEqual(
154       "/",
155       gettempdir([{fun gettempdir_identity/1, ["/--not-here--/"]},
156                   {fun gettempdir_identity/1, ["/"]}],
157                  fun normalize_dir/1)),
158    ?assertEqual(
159       "/",
160       %% simulate a true os:getenv unset env
161       gettempdir([{fun gettempdir_identity/1, [false]},
162                   {fun gettempdir_identity/1, ["/"]}],
163                  fun normalize_dir/1)),
164    ok.
165
166gettempdir_identity_test() ->
167    ?assertEqual(
168       "/",
169       gettempdir([{fun gettempdir_identity/1, ["/"]}], fun normalize_dir/1)),
170    ok.
171
172gettempdir_cwd_test() ->
173    {ok, Cwd} = file:get_cwd(),
174    ?assertEqual(
175       normalize_dir(Cwd),
176       gettempdir([{fun gettempdir_cwd/1, [cwd]}], fun normalize_dir/1)),
177    ok.
178
179rngchars_test() ->
180    crypto:start(),
181    ?assertEqual(
182       "",
183       rngchars(0)),
184    ?assertEqual(
185       10,
186       length(rngchars(10))),
187    ok.
188
189rngchar_test() ->
190    ?assertEqual(
191       $a,
192       rngchar(0)),
193    ?assertEqual(
194       $A,
195       rngchar(26)),
196    ?assertEqual(
197       $_,
198       rngchar(62)),
199    ok.
200
201mkdtemp_n_failonce_test() ->
202    crypto:start(),
203    D = mkdtemp(),
204    Path = filename:join([D, "testdir"]),
205    %% Toggle the existence of a dir so that it fails
206    %% the first time and succeeds the second.
207    F = fun () ->
208                case filelib:is_dir(Path) of
209                    true ->
210                        file:del_dir(Path);
211                    false ->
212                        file:make_dir(Path)
213                end,
214                Path
215        end,
216    try
217        %% Fails the first time
218        ?assertThrow(
219           {error, eexist},
220           mkdtemp_n(F, 1)),
221        %% Reset state
222        file:del_dir(Path),
223        %% Succeeds the second time
224        ?assertEqual(
225           Path,
226           mkdtemp_n(F, 2))
227    after rmtempdir(D)
228    end,
229    ok.
230
231mkdtemp_n_fail_test() ->
232    {ok, Cwd} = file:get_cwd(),
233    ?assertThrow(
234       {error, eexist},
235       mkdtemp_n(fun () -> Cwd end, 1)),
236    ?assertThrow(
237       {error, eexist},
238       mkdtemp_n(fun () -> Cwd end, 2)),
239    ok.
240
241make_dir_fail_test() ->
242    {ok, Cwd} = file:get_cwd(),
243    ?assertThrow(
244      {error, eexist},
245      make_dir(Cwd)),
246    ok.
247
248mkdtemp_test() ->
249    crypto:start(),
250    D = mkdtemp(),
251    ?assertEqual(
252       true,
253       filelib:is_dir(D)),
254    ?assertEqual(
255       ok,
256       file:del_dir(D)),
257    ok.
258
259rmtempdir_test() ->
260    crypto:start(),
261    D1 = mkdtemp(),
262    ?assertEqual(
263       true,
264       filelib:is_dir(D1)),
265    ?assertEqual(
266       ok,
267       rmtempdir(D1)),
268    D2 = mkdtemp(),
269    ?assertEqual(
270       true,
271       filelib:is_dir(D2)),
272    ok = file:write_file(filename:join([D2, "foo"]), <<"bytes">>),
273    D3 = mkdtemp("suffix", "prefix", D2),
274    ?assertEqual(
275       true,
276       filelib:is_dir(D3)),
277    ok = file:write_file(filename:join([D3, "foo"]), <<"bytes">>),
278    ?assertEqual(
279       ok,
280       rmtempdir(D2)),
281    ?assertEqual(
282       {error, enoent},
283       file:consult(D3)),
284    ?assertEqual(
285       {error, enoent},
286       file:consult(D2)),
287    ok.
288
289gettempdir_env_test() ->
290    Env = pushenv(["TMPDIR", "TEMP", "TMP"]),
291    FalseEnv = [{"TMPDIR", false}, {"TEMP", false}, {"TMP", false}],
292    try
293        popenv(FalseEnv),
294        popenv([{"TMPDIR", "/"}]),
295        ?assertEqual(
296           "/",
297           os:getenv("TMPDIR")),
298        ?assertEqual(
299           "/",
300           gettempdir()),
301        {ok, Cwd} = file:get_cwd(),
302        popenv(FalseEnv),
303        popenv([{"TMP", Cwd}]),
304        ?assertEqual(
305           normalize_dir(Cwd),
306           gettempdir())
307    after popenv(Env)
308    end,
309    ok.
310
311-endif.
312