A module $J$ is injective iff every short exact sequence of the form $0\to J\to A\to B\to 0$ splits.
I have seen these similar questions 1, 2, 3, but none contain a proof of this statement above.
Here is my definition of injective module: An ($R$-)module $J$ is injective if for any injection $A\xrightarrow{f} B$ of modules and any map $A\xrightarrow{g} J$, there exists a map $B\xrightarrow{h} J$ such that $h\circ f = g$.
For the forward direction of this statement, it is easy enough. Suppose $J$ is injective. Then for any injection $J \xrightarrow{f} A$, there exists a map $A\xrightarrow{g} J$ such that $g \circ f = \operatorname{id}_J$. So by the splitting lemma, the short exact sequence splits.
However, the reverse direction seems more complicated. We assume that every such sequence splits, and need to show that $J$ is injective. I think this should require more work than a simple application of the definition. Namely, we start with an arbitrary injection $A\xrightarrow{f}B$ and map $A\xrightarrow{g}J$ and need to come up with a map $B\xrightarrow{h} J$ such that $h\circ f = g$. This means we somehow need to convert these maps into a short exact sequence of the given form, and then maybe use the splitting lemma to get the desired map.
I've thought about this for a while, but cannot see how to get a short exact sequence of the above form out of this information. What is the key idea for this direction of the proof?